Improve connection management-related logic

This commit is contained in:
codeking 2025-12-30 01:40:42 +01:00
parent fe8903a807
commit f32d0a8986
5 changed files with 64 additions and 127 deletions

View file

@ -2,7 +2,6 @@ from collections.abc import Callable
from core.Constants import Constants from core.Constants import Constants
from core.Errors import InvalidSubscriptionError, MissingSubscriptionError, ConnectionUnprotectedError, ConnectionTerminationError, CommandNotFoundError from core.Errors import InvalidSubscriptionError, MissingSubscriptionError, ConnectionUnprotectedError, ConnectionTerminationError, CommandNotFoundError
from core.controllers.ConfigurationController import ConfigurationController from core.controllers.ConfigurationController import ConfigurationController
from core.controllers.PolicyController import PolicyController
from core.controllers.ProfileController import ProfileController from core.controllers.ProfileController import ProfileController
from core.controllers.SessionStateController import SessionStateController from core.controllers.SessionStateController import SessionStateController
from core.controllers.SystemStateController import SystemStateController from core.controllers.SystemStateController import SystemStateController
@ -78,10 +77,8 @@ class ConnectionController:
if ConnectionController.system_uses_wireguard_interface() and SystemStateController.exists(): if ConnectionController.system_uses_wireguard_interface() and SystemStateController.exists():
active_profile = ProfileController.get(SystemStateController.get().profile_id)
try: try:
ConnectionController.terminate_system_connection(active_profile) ConnectionController.terminate_system_connection()
except ConnectionTerminationError: except ConnectionTerminationError:
pass pass
@ -125,6 +122,7 @@ class ConnectionController:
session_directory = tempfile.mkdtemp(prefix='hv-') session_directory = tempfile.mkdtemp(prefix='hv-')
session_state = SessionStateController.get_or_new(profile.id) session_state = SessionStateController.get_or_new(profile.id)
maximum_number_of_attempts = None
port_number = None port_number = None
proxy_port_number = None proxy_port_number = None
@ -139,6 +137,7 @@ class ConnectionController:
if profile.connection.code == 'tor': if profile.connection.code == 'tor':
maximum_number_of_attempts = 5
port_number = ConnectionController.get_random_available_port_number() port_number = ConnectionController.get_random_available_port_number()
ConnectionController.establish_tor_session_connection(session_directory, port_number) ConnectionController.establish_tor_session_connection(session_directory, port_number)
session_state.network_port_numbers.append(port_number) session_state.network_port_numbers.append(port_number)
@ -154,6 +153,8 @@ class ConnectionController:
if profile.connection.masked: if profile.connection.masked:
maximum_number_of_attempts = 5
while proxy_port_number is None or proxy_port_number == port_number: while proxy_port_number is None or proxy_port_number == port_number:
proxy_port_number = ConnectionController.get_random_available_port_number() proxy_port_number = ConnectionController.get_random_available_port_number()
@ -161,7 +162,7 @@ class ConnectionController:
session_state.network_port_numbers.append(proxy_port_number) session_state.network_port_numbers.append(proxy_port_number)
if not profile.connection.is_unprotected(): if not profile.connection.is_unprotected():
ConnectionController.await_connection(proxy_port_number or port_number, connection_observer=connection_observer) ConnectionController.await_connection(proxy_port_number or port_number, maximum_number_of_attempts, connection_observer=connection_observer)
SessionStateController.update_or_create(session_state) SessionStateController.update_or_create(session_state)
@ -179,7 +180,7 @@ class ConnectionController:
except ConnectionError: except ConnectionError:
try: try:
ConnectionController.terminate_system_connection(profile) ConnectionController.terminate_system_connection()
except ConnectionTerminationError: except ConnectionTerminationError:
pass pass
@ -188,7 +189,7 @@ class ConnectionController:
except CalledProcessError: except CalledProcessError:
try: try:
ConnectionController.terminate_system_connection(profile) ConnectionController.terminate_system_connection()
except ConnectionTerminationError: except ConnectionTerminationError:
pass pass
@ -198,7 +199,7 @@ class ConnectionController:
except (ConnectionError, CalledProcessError): except (ConnectionError, CalledProcessError):
try: try:
ConnectionController.terminate_system_connection(profile) ConnectionController.terminate_system_connection()
except ConnectionTerminationError: except ConnectionTerminationError:
pass pass
@ -277,26 +278,23 @@ class ConnectionController:
return subprocess.Popen(('proxychains4', '-f', proxychains_configuration_file_path, 'microsocks', '-p', str(proxy_port_number)), stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT) return subprocess.Popen(('proxychains4', '-f', proxychains_configuration_file_path, 'microsocks', '-p', str(proxy_port_number)), stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
@staticmethod @staticmethod
def terminate_system_connection(profile: SystemProfile): def terminate_system_connection():
if shutil.which('pkexec') is None: if shutil.which('nmcli') is None:
raise CommandNotFoundError('pkexec') raise CommandNotFoundError('nmcli')
if shutil.which('wg-quick') is None: if SystemStateController.exists():
raise CommandNotFoundError('wg-quick')
process = subprocess.Popen(('pkexec', 'wg-quick', 'down', profile.get_wireguard_configuration_path()), stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT) process = subprocess.Popen(('nmcli', 'connection', 'delete', 'wg'), stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
completed_successfully = not bool(os.waitpid(process.pid, 0)[1] >> 8) completed_successfully = not bool(os.waitpid(process.pid, 0)[1] >> 8)
if completed_successfully or not ConnectionController.system_uses_wireguard_interface(): if completed_successfully or not ConnectionController.system_uses_wireguard_interface():
system_state = SystemStateController.get() subprocess.run(('nmcli', 'connection', 'delete', 'hv-ipv6-sink'), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
SystemState.dissolve()
if system_state is not None: else:
system_state.dissolve() raise ConnectionTerminationError('The connection could not be terminated.')
else:
raise ConnectionTerminationError('The connection could not be terminated.')
@staticmethod @staticmethod
def get_proxies(port_number: int): def get_proxies(port_number: int):
@ -317,12 +315,11 @@ class ConnectionController:
return port_number return port_number
@staticmethod @staticmethod
def await_connection(port_number: int = None, connection_observer: Optional[ConnectionObserver] = None): def await_connection(port_number: int = None, maximum_number_of_attempts: int = 2, connection_observer: Optional[ConnectionObserver] = None):
if port_number is None: if port_number is None:
ConnectionController.await_network_interface() ConnectionController.await_network_interface()
maximum_number_of_attempts = 5
retry_interval = 5.0 retry_interval = 5.0
for retry_count in range(maximum_number_of_attempts): for retry_count in range(maximum_number_of_attempts):
@ -380,44 +377,34 @@ class ConnectionController:
@staticmethod @staticmethod
def __establish_system_connection(profile: SystemProfile, connection_observer: Optional[ConnectionObserver] = None): def __establish_system_connection(profile: SystemProfile, connection_observer: Optional[ConnectionObserver] = None):
if shutil.which('wg-quick') is None: if shutil.which('nmcli') is None:
raise CommandNotFoundError('wg-quick') raise CommandNotFoundError('nmcli')
privilege_policy = PolicyController.get('privilege') ConnectionController.terminate_system_connection()
permission_denied = False try:
return_code = None process_output = subprocess.check_output(('nmcli', 'connection', 'import', '--temporary', 'type', 'wireguard', 'file', profile.get_wireguard_configuration_path()), text=True)
except CalledProcessError as exception:
raise CalledProcessError(exception.returncode, 'nmcli')
if privilege_policy.is_instated(): connection_id = (m := re.search(r'(?<=\()([a-f0-9-]+?)(?=\))', process_output)) and m.group(1)
subprocess.run(('nmcli', 'connection', 'modify', connection_id, 'ipv4.dns-priority', '-1750', 'ipv4.ignore-auto-dns', 'yes'), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
process = subprocess.Popen(('sudo', '-n', 'wg-quick', 'up', profile.get_wireguard_configuration_path()), stdout=subprocess.DEVNULL, stderr=subprocess.PIPE) try:
process.wait() ipv6_method = subprocess.check_output(('nmcli', '-g', 'ipv6.method', 'connection', 'show', connection_id), text=True).strip()
except CalledProcessError:
raise ConnectionError('The connection could not be established.')
return_code = process.returncode if ipv6_method in ('disabled', 'ignore'):
permission_denied = return_code != 0 and b'sudo:' in process.stderr.read() subprocess.run(('nmcli', 'connection', 'add', 'type', 'dummy', 'save', 'no', 'con-name', 'hv-ipv6-sink', 'ifname', 'hvipv6sink0', 'ipv6.addresses', 'fd7a:fd4b:54e3:077c::/64', 'ipv6.gateway', 'fd7a:fd4b:54e3:077c::1', 'ipv6.route-metric', '72'), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
if not privilege_policy.is_instated() or permission_denied: SystemStateController.create(profile.id)
if shutil.which('pkexec') is None: try:
raise CommandNotFoundError('pkexec') ConnectionController.await_connection(connection_observer=connection_observer)
process = subprocess.Popen(('pkexec', 'wg-quick', 'up', profile.get_wireguard_configuration_path()), stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT) except ConnectionError:
process.wait() raise ConnectionError('The connection could not be established.')
return_code = process.returncode
if return_code == 0:
SystemStateController.update_or_create(SystemState(profile.id))
try:
ConnectionController.await_connection(connection_observer=connection_observer)
except ConnectionError:
raise ConnectionError('The connection could not be established.')
else:
raise CalledProcessError(return_code, 'wg-quick')
@staticmethod @staticmethod
def __with_tor_connection(*args, task: Callable[..., Any], connection_observer: Optional[ConnectionObserver] = None, **kwargs): def __with_tor_connection(*args, task: Callable[..., Any], connection_observer: Optional[ConnectionObserver] = None, **kwargs):
@ -426,7 +413,7 @@ class ConnectionController:
port_number = ConnectionController.get_random_available_port_number() port_number = ConnectionController.get_random_available_port_number()
process = ConnectionController.establish_tor_session_connection(session_directory, port_number) process = ConnectionController.establish_tor_session_connection(session_directory, port_number)
ConnectionController.await_connection(port_number, connection_observer=connection_observer) ConnectionController.await_connection(port_number, 5, connection_observer=connection_observer)
task_output = task(*args, proxies=ConnectionController.get_proxies(port_number), **kwargs) task_output = task(*args, proxies=ConnectionController.get_proxies(port_number), **kwargs)
process.terminate() process.terminate()

View file

@ -112,8 +112,13 @@ class ProfileController:
if SystemStateController.exists(): if SystemStateController.exists():
system_state = SystemStateController.get()
if profile.id != system_state.profile_id:
raise ProfileDeactivationError('The profile could not be disabled.')
try: try:
ConnectionController.terminate_system_connection(profile) ConnectionController.terminate_system_connection()
except ConnectionTerminationError: except ConnectionTerminationError:
raise ProfileDeactivationError('The profile could not be disabled.') raise ProfileDeactivationError('The profile could not be disabled.')

View file

@ -7,20 +7,18 @@ class SystemStateController:
def get(): def get():
return SystemState.get() return SystemState.get()
@staticmethod
def get_or_new(profile_id: int):
system_state = SystemStateController.get()
if system_state is None:
return SystemState(profile_id)
return system_state
@staticmethod @staticmethod
def exists(): def exists():
return SystemState.exists() return SystemState.exists()
@staticmethod
def create(profile_id):
return SystemState(profile_id).save()
@staticmethod @staticmethod
def update_or_create(system_state): def update_or_create(system_state):
system_state.save() system_state.save()
@staticmethod
def dissolve():
return SystemState.dissolve()

View file

@ -1,12 +1,8 @@
from core.Constants import Constants from core.Constants import Constants
from core.Errors import CommandNotFoundError, PolicyInstatementError, PolicyRevocationError, PolicyAssignmentError from core.Errors import CommandNotFoundError, PolicyInstatementError, PolicyRevocationError, PolicyAssignmentError
from core.models.BasePolicy import BasePolicy from core.models.BasePolicy import BasePolicy
from packaging import version
from packaging.version import InvalidVersion
from subprocess import CalledProcessError
import os import os
import pwd import pwd
import re
import shutil import shutil
import subprocess import subprocess
@ -19,33 +15,7 @@ class PrivilegePolicy(BasePolicy):
return self.__generate(username) return self.__generate(username)
def instate(self): def instate(self):
pass
if shutil.which('pkexec') is None:
raise CommandNotFoundError('pkexec')
if not self.__is_compatible():
raise PolicyInstatementError('The privilege policy is not compatible.')
username = self.__determine_username()
privilege_policy = self.__generate(username)
completed_successfully = False
failed_attempt_count = 0
while not completed_successfully and failed_attempt_count < 3:
process = subprocess.Popen((
'pkexec', 'install', '/dev/stdin', Constants.HV_PRIVILEGE_POLICY_PATH, '-o', 'root', '-m', '440'
), stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
process.communicate(f'{privilege_policy}\n')
completed_successfully = (process.returncode == 0)
if not completed_successfully:
failed_attempt_count += 1
if not completed_successfully:
raise PolicyInstatementError('The privilege policy could not be instated.')
def revoke(self): def revoke(self):
@ -86,25 +56,4 @@ class PrivilegePolicy(BasePolicy):
@staticmethod @staticmethod
def __is_compatible(): def __is_compatible():
return False
try:
process_output = subprocess.check_output(('sudo', '-V'), text=True)
except (CalledProcessError, FileNotFoundError):
return False
if process_output.splitlines():
sudo_version_details = process_output.splitlines()[0].strip()
else:
return False
sudo_version_number = (m := re.search(r'(\d[0-9.]+?)(?=p|$)', sudo_version_details)) and m.group(1)
if not sudo_version_number:
return False
try:
sudo_version = version.parse(sudo_version_number)
except InvalidVersion:
return False
return sudo_version >= version.parse('1.9.10') and os.path.isfile('/usr/bin/wg-quick')

View file

@ -30,17 +30,15 @@ class SystemState:
def get(): def get():
try: try:
system_state_file_contents = open(f'{SystemState.__get_state_path()}/system.json', 'r').read() system_state_file_contents = open(f'{SystemState.__get_state_path()}/system.json', 'r').read()
except FileNotFoundError: system_state_dict = json.loads(system_state_file_contents)
return None
try: # noinspection PyUnresolvedReferences
system_state = json.loads(system_state_file_contents) return SystemState.from_dict(system_state_dict)
except JSONDecodeError:
return None
# noinspection PyUnresolvedReferences except (FileNotFoundError, JSONDecodeError, KeyError):
return SystemState.from_dict(system_state) return None
@staticmethod @staticmethod
def exists(): def exists():