State + FSM + Exports: Implement changing a secure internet location

5 files changed, 43 insertions, 4 deletions

diff --git a/exports/exports.go b/exports/exports.go
index 1e17587..2eb3ab7 100644
--- a/exports/exports.go
+++ b/exports/exports.go
@@ -193,6 +193,17 @@ func SetProfileID(name *C.char, data *C.char) *C.char {
 	return C.CString(ErrorToString(profileErr))
 }
 
+//export ChangeSecureLocation
+func ChangeSecureLocation(name *C.char) *C.char {
+	nameStr := C.GoString(name)
+	state, stateErr := GetVPNState(nameStr)
+	if stateErr != nil {
+		return C.CString(ErrorToString(stateErr))
+	}
+	locationErr := state.ChangeSecureLocation()
+	return C.CString(ErrorToString(locationErr))
+}
+
 //export SetSecureLocation
 func SetSecureLocation(name *C.char, data *C.char) *C.char {
 	nameStr := C.GoString(name)
diff --git a/internal/fsm/fsm.go b/internal/fsm/fsm.go
index b822a72..8b5875b 100644
--- a/internal/fsm/fsm.go
+++ b/internal/fsm/fsm.go
@@ -132,9 +132,9 @@ type FSM struct {
 func (fsm *FSM) Init(name string, callback func(string, string, string), directory string, debug bool) {
 	fsm.States = FSMStates{
 		DEREGISTERED:   FSMState{Transitions: []FSMTransition{{NO_SERVER, "Client registers"}}},
-		NO_SERVER:      FSMState{Transitions: []FSMTransition{{CHOSEN_SERVER, "User chooses a server"}, {SEARCH_SERVER, "The user is trying to choose a Server in the UI"}, {CONNECTED, "The user is already connected"}}},
+		NO_SERVER:      FSMState{Transitions: []FSMTransition{{CHOSEN_SERVER, "User chooses a server"}, {SEARCH_SERVER, "The user is trying to choose a Server in the UI"}, {CONNECTED, "The user is already connected"}, {ASK_LOCATION, "Change the location in the main screen"}}},
 		SEARCH_SERVER:  FSMState{Transitions: []FSMTransition{{LOADING_SERVER, "User clicks a server in the UI"}, {NO_SERVER, "Cancel or Error"}}, BackState: NO_SERVER},
-		ASK_LOCATION:   FSMState{Transitions: []FSMTransition{{CHOSEN_SERVER, "Location chosen"}, {NO_SERVER, "Cancel or Error"}, {SEARCH_SERVER, "Cancel or Error"}}},
+		ASK_LOCATION:   FSMState{Transitions: []FSMTransition{{CHOSEN_SERVER, "Location chosen"}, {NO_SERVER, "Go back or Error"}, {SEARCH_SERVER, "Cancel or Error"}}},
 		LOADING_SERVER: FSMState{Transitions: []FSMTransition{{CHOSEN_SERVER, "Server info loaded"}, {ASK_LOCATION, "User chooses a Secure Internet server but no location is configured"}}},
 		CHOSEN_SERVER:  FSMState{Transitions: []FSMTransition{{AUTHORIZED, "Found tokens in config"}, {OAUTH_STARTED, "No tokens found in config"}}},
 		OAUTH_STARTED:  FSMState{Transitions: []FSMTransition{{AUTHORIZED, "User authorizes with browser"}, {NO_SERVER, "Cancel or Error"}, {SEARCH_SERVER, "Cancel or Error"}}, BackState: NO_SERVER},
diff --git a/state.go b/state.go
index 2f7e785..d559c88 100644
--- a/state.go
+++ b/state.go
@@ -154,7 +154,7 @@ func (state *VPNState) SetSecureLocation(countryCode string) error {
 	return nil
 }
 
-func (state *VPNState) AskSecureLocation() error {
+func (state *VPNState) askSecureLocation() error {
 	errorMessage := "failed asking Secure Internet location"
 	locations, locationsErr := state.Discovery.GetSecureLocationList()
 	if locationsErr != nil {
@@ -184,7 +184,7 @@ func (state *VPNState) addSecureInternetHomeServer(orgID string) (server.Server,
 	var locationErr error
 
 	if !state.Servers.HasSecureLocation() {
-		locationErr = state.AskSecureLocation()
+		locationErr = state.askSecureLocation()
 	} else {
 		// reinitialize
 		locationErr = state.SetSecureLocation(state.Servers.GetSecureLocation())
@@ -285,6 +285,25 @@ func (state *VPNState) CancelOAuth() error {
 	return nil
 }
 
+func (state *VPNState) ChangeSecureLocation() error {
+	errorMessage := "failed to change location from the main screen"
+
+	if !state.FSM.InState(fsm.NO_SERVER) {
+		return &types.WrappedErrorMessage{Message: errorMessage, Err: fsm.WrongStateError{Got: state.FSM.Current, Want: fsm.NO_SERVER}.CustomError()}
+	}
+
+	askLocationErr := state.askSecureLocation()
+
+	if askLocationErr != nil {
+		return &types.WrappedErrorMessage{Message: errorMessage, Err: askLocationErr}
+	}
+
+	// Go back to the main screen
+	state.FSM.GoTransitionWithData(fsm.NO_SERVER, state.GetSavedServers(), false)
+
+	return nil
+}
+
 func (state *VPNState) GetDiscoOrganizations() (string, error) {
 	if state.FSM.InState(fsm.DEREGISTERED) {
 		return "", &types.WrappedErrorMessage{Message: "failed to get the organizations with Discovery", Err: fsm.DeregisteredError{}.CustomError()}
diff --git a/wrappers/python/src/__init__.py b/wrappers/python/src/__init__.py
index ed6fcc7..1e42bb5 100644
--- a/wrappers/python/src/__init__.py
+++ b/wrappers/python/src/__init__.py
@@ -77,6 +77,7 @@ lib.GetDiscoServers.argtypes, lib.GetDiscoServers.restype = [c_char_p], DataErro
 lib.GoBack.argtypes, lib.GoBack.restype = [c_char_p], None
 lib.CancelOAuth.argtypes, lib.CancelOAuth.restype = [c_char_p], c_void_p
 lib.SetProfileID.argtypes, lib.SetProfileID.restype = [c_char_p, c_char_p], c_void_p
+lib.ChangeSecureLocation.argtypes, lib.ChangeSecureLocation.restype = [c_char_p], c_void_p
 lib.SetSecureLocation.argtypes, lib.SetSecureLocation.restype = [
     c_char_p,
     c_char_p,
diff --git a/wrappers/python/src/main.py b/wrappers/python/src/main.py
index 87d2b53..784e622 100644
--- a/wrappers/python/src/main.py
+++ b/wrappers/python/src/main.py
@@ -187,6 +187,14 @@ class EduVPN(object):
         if profile_err:
             raise Exception(profile_err)
 
+    def change_secure_location(self) -> None:
+        # Set the location by country code
+        self.location_event = threading.Event()
+        location_err = self.go_function(lib.ChangeSecureLocation)
+
+        if location_err:
+            raise Exception(location_err)
+
     def set_secure_location(self, country_code: str) -> None:
         # Set the location by country code
         location_err = self.go_function(lib.SetSecureLocation, country_code)