Secure Internet: Implement the Ask Location transition callback

3 files changed, 36 insertions, 6 deletions

diff --git a/wrappers/python/main.py b/wrappers/python/main.py
index 5422d93..d75504f 100644
--- a/wrappers/python/main.py
+++ b/wrappers/python/main.py
@@ -1,6 +1,7 @@
 import eduvpncommon.main as eduvpn
 import webbrowser
 import json
+import sys
 
 # Asks the user for a profile index
 # It loops up until a valid input is given
@@ -31,6 +32,11 @@ def setup_callbacks(_eduvpn: eduvpn.EduVPN) -> None:
         print(f"Got OAuth URL {url}, old state: {old_state}")
         webbrowser.open(url)
 
+    @_eduvpn.event.on("Ask_Location", eduvpn.StateType.Enter)
+    def ask_location(old_state: str, locations: str):
+        print("Locations: ", locations)
+        _eduvpn.set_secure_location("NL")
+
     # The callback which asks the user for a profile
     @_eduvpn.event.on("Ask_Profile", eduvpn.StateType.Enter)
     def ask_profile(old_state: str, profiles: str):
@@ -79,9 +85,12 @@ if __name__ == "__main__":
     # Get a Wireguard/OpenVPN config
     try:
         config, config_type = _eduvpn.get_config_custom_server(server)
+        print(f"Got a config with type: {config_type} and contents:\n{config}")
     except Exception as e:
         print("Failed to connect:", e)
-    print(f"Got a config with type: {config_type} and contents:\n{config}")
+        # Save and exit
+        _eduvpn.deregister()
+        sys.exit(1)
 
     # Set the internal FSM state to connected
     try:
diff --git a/wrappers/python/src/__init__.py b/wrappers/python/src/__init__.py
index f2ae66e..22d5406 100644
--- a/wrappers/python/src/__init__.py
+++ b/wrappers/python/src/__init__.py
@@ -75,6 +75,7 @@ lib.GetOrganizationsList.argtypes, lib.GetOrganizationsList.restype = [
 lib.GetServersList.argtypes, lib.GetServersList.restype = [c_char_p], DataError
 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.SetSecureLocation.argtypes, lib.SetSecureLocation.restype = [c_char_p, c_char_p], c_void_p
 lib.SetConnected.argtypes, lib.SetConnected.restype = [c_char_p], c_void_p
 lib.SetDisconnected.argtypes, lib.SetDisconnected.restype = [c_char_p], c_void_p
 lib.GetIdentifier.argtypes, lib.GetIdentifier.restype = [c_char_p], DataError
diff --git a/wrappers/python/src/main.py b/wrappers/python/src/main.py
index dda3250..542fc54 100644
--- a/wrappers/python/src/main.py
+++ b/wrappers/python/src/main.py
@@ -38,12 +38,18 @@ class EduVPN(object):
         # The ask profile callback needs to wait for the UI thread to select a profile
         # This is stored in the profile_event
         self.profile_event: Optional[threading.Event] = None
+        self.location_event: Optional[threading.Event] = None
 
         @self.event.on("Ask_Profile", StateType.Wait)
         def wait_profile_event(old_state: str, profiles: str):
             if self.profile_event:
                 self.profile_event.wait()
 
+        @self.event.on("Ask_Location", StateType.Wait)
+        def wait_location_event(old_state: str, locations: str):
+            if self.location_event:
+                self.location_event.wait()
+
     def go_function(self, func, *args):
         # The functions all have at least one arg type which is the name of the client
         args_gen = encode_args(list(args), func.argtypes[1:])
@@ -102,11 +108,12 @@ class EduVPN(object):
             func, url, force_tcp
         )
 
+        self.profile_event = None
+        self.location_event = None
+
         if config_err:
             raise Exception(config_err)
 
-        self.profile_event = None
-
         return config, config_type
 
     def get_config_custom_server(
@@ -122,6 +129,7 @@ class EduVPN(object):
     def get_config_secure_internet(
         self, url: str, force_tcp: bool = False
     ) -> Tuple[str, str]:
+        self.location_event = threading.Event()
         return self.get_config(url, lib.GetConfigSecureInternet, force_tcp)
 
     def set_connected(self) -> None:
@@ -173,10 +181,22 @@ class EduVPN(object):
         # Set the profile id
         profile_err = self.go_function(lib.SetProfileID, profile_id)
 
-        if profile_err:
-            raise Exception(profile_err)
-
         # If there is a profile event, set it so that the wait callback finishes
         # And so that the Go code can move to the next state
         if self.profile_event:
             self.profile_event.set()
+
+        if profile_err:
+            raise Exception(profile_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)
+
+        # If there is a location event, set it so that the wait callback finishes
+        # And so that the Go code can move to the next state
+        if self.location_event:
+            self.location_event.set()
+
+        if location_err:
+            raise Exception(location_err)