From 2683a1655afd3d10d9dfca6c5c3fb0f19deabb0b Mon Sep 17 00:00:00 2001 From: codeking Date: Thu, 19 Sep 2024 01:56:19 +0200 Subject: [PATCH] Update and refactor connection-related logic --- core/Errors.py | 4 ++++ core/controllers/ConnectionController.py | 30 ++++++++++++++++++++++-- core/controllers/ProfileController.py | 26 ++++++++++---------- 3 files changed, 44 insertions(+), 16 deletions(-) diff --git a/core/Errors.py b/core/Errors.py index f492a8a..a3cb5cb 100644 --- a/core/Errors.py +++ b/core/Errors.py @@ -2,6 +2,10 @@ class UnknownConnectionTypeError(Exception): pass +class ConnectionTerminationError(Exception): + pass + + class ProfileStateConflictError(Exception): pass diff --git a/core/controllers/ConnectionController.py b/core/controllers/ConnectionController.py index 35c16b5..0809039 100644 --- a/core/controllers/ConnectionController.py +++ b/core/controllers/ConnectionController.py @@ -1,5 +1,5 @@ from core.Constants import Constants -from core.Errors import InvalidSubscriptionError, MissingSubscriptionError, ConnectionUnprotectedError, ProfileActivationError +from core.Errors import InvalidSubscriptionError, MissingSubscriptionError, ConnectionUnprotectedError, ConnectionTerminationError from core.controllers.ConfigurationController import ConfigurationController from core.controllers.ProfileController import ProfileController from core.controllers.SessionStateController import SessionStateController @@ -139,6 +139,12 @@ class ConnectionController: @staticmethod def establish_system_connection(profile: SystemProfile): + if subprocess.getstatusoutput('pkexec --help')[0] == 127: + raise OSError('The polkit toolkit does not appear to be installed.') + + if subprocess.getstatusoutput('wg-quick --help')[0] == 127: + raise OSError('WireGuard tools does not appear to be installed.') + 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) @@ -158,7 +164,7 @@ class ConnectionController: SystemStateController.update_or_create(SystemState(profile.id)) else: - raise ProfileActivationError('The profile could not be enabled.') + raise ConnectionError('The connection could not be established.') time.sleep(1.0) @@ -227,6 +233,26 @@ class ConnectionController: return subprocess.Popen(('proxychains4', '-f', proxychains_configuration_file_path, 'microsocks', '-p', str(proxy_port_number)), stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT) + @staticmethod + def terminate_system_connection(profile: SystemProfile): + + if subprocess.getstatusoutput('pkexec --help')[0] == 127: + raise OSError('The polkit toolkit does not appear to be installed.') + + if subprocess.getstatusoutput('wg-quick --help')[0] == 127: + raise OSError('WireGuard tools does not appear to be installed.') + + process = subprocess.Popen(('pkexec', 'wg-quick', 'down', 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 or not ConnectionController.system_uses_wireguard_interface(): + + system_state = SystemStateController.get() + system_state.dissolve() + + else: + raise ConnectionTerminationError('The connection could not be terminated.') + @staticmethod def get_proxies(port_number: int): diff --git a/core/controllers/ProfileController.py b/core/controllers/ProfileController.py index a818e30..171505f 100644 --- a/core/controllers/ProfileController.py +++ b/core/controllers/ProfileController.py @@ -1,4 +1,4 @@ -from core.Errors import ProfileStateConflictError, InvalidSubscriptionError, MissingSubscriptionError, ConnectionUnprotectedError, ProfileDeactivationError +from core.Errors import ProfileStateConflictError, InvalidSubscriptionError, MissingSubscriptionError, ConnectionUnprotectedError, ConnectionTerminationError, ProfileActivationError, ProfileDeactivationError from core.controllers.ApplicationController import ApplicationController from core.controllers.ApplicationVersionController import ApplicationVersionController from core.controllers.SessionStateController import SessionStateController @@ -12,8 +12,6 @@ from core.observers.ConnectionObserver import ConnectionObserver from core.observers.ProfileObserver import ProfileObserver from core.services.WebServiceApiService import WebServiceApiService from typing import Union, Optional -import os -import subprocess import time @@ -54,7 +52,10 @@ class ProfileController: if not application_version.is_installed(): ApplicationVersionController.install(application_version, application_version_observer=application_version_observer, connection_observer=connection_observer) - port_number = ConnectionController.establish_connection(profile, force=force, connection_observer=connection_observer) + try: + port_number = ConnectionController.establish_connection(profile, force=force, connection_observer=connection_observer) + except ConnectionError: + raise ProfileActivationError('The profile could not be enabled.') if profile_observer is not None: profile_observer.notify('enabled', profile) @@ -63,7 +64,10 @@ class ProfileController: if profile.is_system_profile(): - ConnectionController.establish_connection(profile, force=force, connection_observer=connection_observer) + try: + ConnectionController.establish_connection(profile, force=force, connection_observer=connection_observer) + except ConnectionError: + raise ProfileActivationError('The profile could not be enabled.') if profile_observer is not None: profile_observer.notify('enabled', profile) @@ -93,15 +97,9 @@ class ProfileController: if SystemStateController.exists(): - process = subprocess.Popen(('pkexec', 'wg-quick', 'down', 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 or not ConnectionController.system_uses_wireguard_interface(): - - system_state = SystemStateController.get() - system_state.dissolve() - - else: + try: + ConnectionController.terminate_system_connection(profile) + except ConnectionTerminationError: raise ProfileDeactivationError('The profile could not be disabled.') if profile_observer is not None: