Implement initial support for session renegotiation

core/Errors.py

@@ -34,6 +34,10 @@ class ApplicationAlreadyInstalledError(Exception):
     pass
 
 
+class MissingLocationError(Exception):
+    pass
+
+
 class MissingSubscriptionError(Exception):
     pass
 

core/controllers/ConnectionController.py

@@ -62,21 +62,52 @@ class ConnectionController:
                 if not profile.subscription.has_been_activated():
                     ProfileController.activate_subscription(profile, connection_observer=connection_observer)
 
-                wireguard_configuration = ConnectionController.with_preferred_connection(profile.location.code, profile.subscription.billing_code, task=WebServiceApiService.post_wireguard_session, connection_observer=connection_observer)
+                ProfileController.register_wireguard_session(profile, connection_observer=connection_observer)
-
-                if wireguard_configuration is None:
-                    raise InvalidSubscriptionError()
-
-                profile.attach_wireguard_configuration(wireguard_configuration)
 
             else:
+
+                if profile.is_system_profile():
+
+                    if ConnectionController.system_uses_wireguard_interface() and SystemStateController.exists():
+
+                        active_profile = ProfileController.get(SystemStateController.get().profile_id)
+
+                        try:
+                            ConnectionController.terminate_system_connection(active_profile)
+                        except ConnectionTerminationError:
+                            pass
+
                 raise MissingSubscriptionError()
 
         if profile.is_session_profile():
+
+            try:
                 return ConnectionController.establish_session_connection(profile, force=force, connection_observer=connection_observer)
 
+            except ConnectionError:
+
+                if ConnectionController.__should_renegotiate(profile):
+
+                    ProfileController.register_wireguard_session(profile, connection_observer=connection_observer)
+                    return ConnectionController.establish_session_connection(profile, force=force, connection_observer=connection_observer)
+
+                else:
+                    raise ConnectionError('The connection could not be established.')
+
         if profile.is_system_profile():
-            return ConnectionController.establish_system_connection(profile)
+
+            try:
+                return ConnectionController.establish_system_connection(profile, connection_observer=connection_observer)
+
+            except ConnectionError:
+
+                if ConnectionController.__should_renegotiate(profile):
+
+                    ProfileController.register_wireguard_session(profile, connection_observer=connection_observer)
+                    return ConnectionController.establish_system_connection(profile, connection_observer=connection_observer)
+
+                else:
+                    raise ConnectionError('The connection could not be established.')
 
     @staticmethod
     def establish_session_connection(profile: SessionProfile, force: Optional[bool] = False, connection_observer: Optional[ConnectionObserver] = None):
@@ -124,7 +155,7 @@ class ConnectionController:
         return proxy_port_number or port_number
 
     @staticmethod
-    def establish_system_connection(profile: SystemProfile):
+    def establish_system_connection(profile: SystemProfile, connection_observer: Optional[ConnectionObserver] = None):
 
         if subprocess.getstatusoutput('pkexec --help')[0] == 127:
             raise OSError('The polkit toolkit does not appear to be installed.')
@@ -136,20 +167,36 @@ class ConnectionController:
         completed_successfully = not bool(os.waitpid(process.pid, 0)[1] >> 8)
 
         if completed_successfully:
+
             SystemStateController.update_or_create(SystemState(profile.id))
 
+            try:
+                ConnectionController.await_connection(connection_observer=connection_observer)
+
+            except ConnectionError:
+
+                ConnectionController.terminate_system_connection(profile)
+                raise ConnectionError('The connection could not be established.')
+
         else:
 
-            ProfileController.disable(profile)
+            ConnectionController.terminate_system_connection(profile)
 
-            process_1 = subprocess.Popen(('pkexec', 'wg-quick', 'down', profile.get_wireguard_configuration_path()), stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
+            process = subprocess.Popen(('pkexec', 'wg-quick', 'up', profile.get_wireguard_configuration_path()), stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
-            process_2 = subprocess.Popen(('pkexec', 'wg-quick', 'up', profile.get_wireguard_configuration_path()), stdin=process_1.stdout, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
+            completed_successfully = not bool(os.waitpid(process.pid, 0)[1] >> 8)
-
-            completed_successfully = not bool(os.waitpid(process_2.pid, 0)[1] >> 8)
 
             if completed_successfully:
+
                 SystemStateController.update_or_create(SystemState(profile.id))
 
+                try:
+                    ConnectionController.await_connection(connection_observer=connection_observer)
+
+                except ConnectionError:
+
+                    ConnectionController.terminate_system_connection(profile)
+                    raise ConnectionError('The connection could not be established.')
+
             else:
                 raise ConnectionError('The connection could not be established.')
 
@@ -240,6 +287,8 @@ class ConnectionController:
         if completed_successfully or not ConnectionController.system_uses_wireguard_interface():
 
             system_state = SystemStateController.get()
+
+            if system_state is not None:
                 system_state.dissolve()
 
         else:
@@ -264,7 +313,10 @@ class ConnectionController:
         return port_number
 
     @staticmethod
-    def await_connection(port_number: int, connection_observer: Optional[ConnectionObserver] = None):
+    def await_connection(port_number: int = None, connection_observer: Optional[ConnectionObserver] = None):
+
+        if port_number is None:
+            ConnectionController.await_network_interface()
 
         maximum_number_of_attempts = 5
         retry_interval = 5.0
@@ -281,7 +333,7 @@ class ConnectionController:
 
             try:
 
-                ConnectionController.__test_proxy_connection(port_number)
+                ConnectionController.__test_connection(port_number)
                 return
 
             except ConnectionError:
@@ -291,6 +343,25 @@ class ConnectionController:
 
         raise ConnectionError('The connection could not be established.')
 
+    @staticmethod
+    def await_network_interface():
+
+        network_interface_is_activated = False
+
+        retry_interval = .5
+        maximum_number_of_attempts = 10
+        attempt = 0
+
+        while not network_interface_is_activated and attempt < maximum_number_of_attempts:
+
+            time.sleep(retry_interval)
+
+            network_interface_is_activated = ConnectionController.system_uses_wireguard_interface()
+            attempt += 1
+
+        if network_interface_is_activated is False:
+            ConnectionError('The network interface could not be activated.')
+
     @staticmethod
     def system_uses_wireguard_interface():
 
@@ -316,11 +387,25 @@ class ConnectionController:
         return task_output
 
     @staticmethod
-    def __test_proxy_connection(port_number: int, timeout: float = 10.0):
+    def __test_connection(port_number: Optional[int] = None, timeout: float = 10.0):
 
         timeout = float(timeout)
+        proxies = None
+
+        if port_number is not None:
+            proxies = ConnectionController.get_proxies(port_number)
 
         try:
-            requests.get(f'{Constants.SP_API_BASE_URL}/health', timeout=timeout, proxies=ConnectionController.get_proxies(port_number))
+            requests.get(f'{Constants.SP_API_BASE_URL}/health', timeout=timeout, proxies=proxies)
         except requests.exceptions.RequestException:
             raise ConnectionError('The connection could not be established.')
+
+    @staticmethod
+    def __should_renegotiate(profile: Union[SessionProfile, SystemProfile]):
+
+        if profile.connection.needs_wireguard_configuration() and profile.has_wireguard_configuration():
+
+            if profile.subscription.has_been_activated():
+                return True
+
+        return False

core/controllers/ProfileController.py

@@ -1,4 +1,4 @@
-from core.Errors import ProfileStateConflictError, InvalidSubscriptionError, MissingSubscriptionError, ConnectionUnprotectedError, ConnectionTerminationError, ProfileActivationError, ProfileDeactivationError
+from core.Errors import ProfileStateConflictError, InvalidSubscriptionError, MissingSubscriptionError, ConnectionUnprotectedError, ConnectionTerminationError, ProfileActivationError, ProfileDeactivationError, MissingLocationError
 from core.controllers.ApplicationController import ApplicationController
 from core.controllers.ApplicationVersionController import ApplicationVersionController
 from core.controllers.SessionStateController import SessionStateController
@@ -211,17 +211,21 @@ class ProfileController:
         profile.has_proxy_configuration()
 
     @staticmethod
-    def register_wireguard_session(profile: Union[SessionProfile, SystemProfile]):
+    def register_wireguard_session(profile: Union[SessionProfile, SystemProfile], connection_observer: Optional[ConnectionObserver] = None):
 
-        if profile.location is None:
+        from core.controllers.ConnectionController import ConnectionController
-            return None
 
         if profile.subscription is None:
-            return None
+            raise InvalidSubscriptionError()
 
-        wireguard_configuration = WebServiceApiService.post_wireguard_session(profile.location.code, profile.subscription.billing_code)
+        if profile.location is None:
+            raise MissingLocationError()
+
+        wireguard_configuration = ConnectionController.with_preferred_connection(profile.location.code, profile.subscription.billing_code, task=WebServiceApiService.post_wireguard_session, connection_observer=connection_observer)
+
+        if wireguard_configuration is None:
+            raise InvalidSubscriptionError()
 
-        if wireguard_configuration is not None:
         profile.attach_wireguard_configuration(wireguard_configuration)
 
     @staticmethod