Update and refactor connection-related logic

This commit is contained in:
codeking 2024-09-19 01:56:19 +02:00
parent 68989396cd
commit 2683a1655a
3 changed files with 44 additions and 16 deletions

View file

@ -2,6 +2,10 @@ class UnknownConnectionTypeError(Exception):
pass
class ConnectionTerminationError(Exception):
pass
class ProfileStateConflictError(Exception):
pass

View file

@ -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):

View file

@ -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: