From f8c92f0f927d64501d388c4d5707d06848f34cd7 Mon Sep 17 00:00:00 2001 From: codeking Date: Sun, 13 Oct 2024 18:38:05 +0200 Subject: [PATCH] Implement initial support for session renegotiation --- core/Errors.py | 4 + core/controllers/ConnectionController.py | 123 +++++++++++++++++++---- core/controllers/ProfileController.py | 20 ++-- 3 files changed, 120 insertions(+), 27 deletions(-) diff --git a/core/Errors.py b/core/Errors.py index eb14ae6..2d46fab 100644 --- a/core/Errors.py +++ b/core/Errors.py @@ -34,6 +34,10 @@ class ApplicationAlreadyInstalledError(Exception): pass +class MissingLocationError(Exception): + pass + + class MissingSubscriptionError(Exception): pass diff --git a/core/controllers/ConnectionController.py b/core/controllers/ConnectionController.py index 91fdc74..4a78efe 100644 --- a/core/controllers/ConnectionController.py +++ b/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) - - if wireguard_configuration is None: - raise InvalidSubscriptionError() - - profile.attach_wireguard_configuration(wireguard_configuration) + ProfileController.register_wireguard_session(profile, connection_observer=connection_observer) 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(): - return ConnectionController.establish_session_connection(profile, force=force, connection_observer=connection_observer) + + 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_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_2.pid, 0)[1] >> 8) + process = subprocess.Popen(('pkexec', 'wg-quick', 'up', profile.get_wireguard_configuration_path()), stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT) + 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: raise ConnectionError('The connection could not be established.') @@ -240,7 +287,9 @@ class ConnectionController: if completed_successfully or not ConnectionController.system_uses_wireguard_interface(): system_state = SystemStateController.get() - system_state.dissolve() + + if system_state is not None: + system_state.dissolve() else: raise ConnectionTerminationError('The connection could not be terminated.') @@ -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 diff --git a/core/controllers/ProfileController.py b/core/controllers/ProfileController.py index 6cf66a3..223bab6 100644 --- a/core/controllers/ProfileController.py +++ b/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,18 +211,22 @@ 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: - return None + from core.controllers.ConnectionController import ConnectionController 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() - if wireguard_configuration is not None: - profile.attach_wireguard_configuration(wireguard_configuration) + 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() + + profile.attach_wireguard_configuration(wireguard_configuration) @staticmethod def get_wireguard_configuration_path(profile: Union[SessionProfile, SystemProfile]):