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.Errors import InvalidSubscriptionError, MissingSubscriptionError, ConnectionUnprotectedError, ConnectionTerminationError, CommandNotFoundError
from core.controllers.ConfigurationController import ConfigurationController
from core.controllers.PolicyController import PolicyController
from core.controllers.ProfileController import ProfileController
from core.controllers.SessionStateController import SessionStateController
from core.controllers.SystemStateController import SystemStateController
@ -78,10 +77,8 @@ class ConnectionController:
if ConnectionController.system_uses_wireguard_interface() and SystemStateController.exists():
active_profile = ProfileController.get(SystemStateController.get().profile_id)
try:
ConnectionController.terminate_system_connection(active_profile)
ConnectionController.terminate_system_connection()
except ConnectionTerminationError:
pass
@ -125,6 +122,7 @@ class ConnectionController:
session_directory = tempfile.mkdtemp(prefix='hv-')
session_state = SessionStateController.get_or_new(profile.id)
maximum_number_of_attempts = None
port_number = None
proxy_port_number = None
@ -139,6 +137,7 @@ class ConnectionController:
if profile.connection.code == 'tor':
maximum_number_of_attempts = 5
port_number = ConnectionController.get_random_available_port_number()
ConnectionController.establish_tor_session_connection(session_directory, port_number)
session_state.network_port_numbers.append(port_number)
@ -154,6 +153,8 @@ class ConnectionController:
if profile.connection.masked:
maximum_number_of_attempts = 5
while proxy_port_number is None or proxy_port_number == 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)
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)
@ -179,7 +180,7 @@ class ConnectionController:
except ConnectionError:
try:
ConnectionController.terminate_system_connection(profile)
ConnectionController.terminate_system_connection()
except ConnectionTerminationError:
pass
@ -188,7 +189,7 @@ class ConnectionController:
except CalledProcessError:
try:
ConnectionController.terminate_system_connection(profile)
ConnectionController.terminate_system_connection()
except ConnectionTerminationError:
pass
@ -198,7 +199,7 @@ class ConnectionController:
except (ConnectionError, CalledProcessError):
try:
ConnectionController.terminate_system_connection(profile)
ConnectionController.terminate_system_connection()
except ConnectionTerminationError:
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)
@staticmethod
def terminate_system_connection(profile: SystemProfile):
def terminate_system_connection():
if shutil.which('pkexec') is None:
raise CommandNotFoundError('pkexec')
if shutil.which('nmcli') is None:
raise CommandNotFoundError('nmcli')
if shutil.which('wg-quick') is None:
raise CommandNotFoundError('wg-quick')
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)
process = subprocess.Popen(('nmcli', 'connection', 'delete', 'wg'), 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():
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:
system_state.dissolve()
else:
raise ConnectionTerminationError('The connection could not be terminated.')
else:
raise ConnectionTerminationError('The connection could not be terminated.')
@staticmethod
def get_proxies(port_number: int):
@ -317,12 +315,11 @@ class ConnectionController:
return port_number
@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:
ConnectionController.await_network_interface()
maximum_number_of_attempts = 5
retry_interval = 5.0
for retry_count in range(maximum_number_of_attempts):
@ -380,44 +377,34 @@ class ConnectionController:
@staticmethod
def __establish_system_connection(profile: SystemProfile, connection_observer: Optional[ConnectionObserver] = None):
if shutil.which('wg-quick') is None:
raise CommandNotFoundError('wg-quick')
if shutil.which('nmcli') is None:
raise CommandNotFoundError('nmcli')
privilege_policy = PolicyController.get('privilege')
ConnectionController.terminate_system_connection()
permission_denied = False
return_code = None
try:
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)
process.wait()
try:
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
permission_denied = return_code != 0 and b'sudo:' in process.stderr.read()
if ipv6_method in ('disabled', 'ignore'):
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:
raise CommandNotFoundError('pkexec')
try:
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)
process.wait()
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')
except ConnectionError:
raise ConnectionError('The connection could not be established.')
@staticmethod
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()
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)
process.terminate()

View file

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

View file

@ -7,20 +7,18 @@ class SystemStateController:
def 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
def exists():
return SystemState.exists()
@staticmethod
def create(profile_id):
return SystemState(profile_id).save()
@staticmethod
def update_or_create(system_state):
system_state.save()
@staticmethod
def dissolve():
return SystemState.dissolve()

View file

@ -1,12 +1,8 @@
from core.Constants import Constants
from core.Errors import CommandNotFoundError, PolicyInstatementError, PolicyRevocationError, PolicyAssignmentError
from core.models.BasePolicy import BasePolicy
from packaging import version
from packaging.version import InvalidVersion
from subprocess import CalledProcessError
import os
import pwd
import re
import shutil
import subprocess
@ -19,33 +15,7 @@ class PrivilegePolicy(BasePolicy):
return self.__generate(username)
def instate(self):
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.')
pass
def revoke(self):
@ -86,25 +56,4 @@ class PrivilegePolicy(BasePolicy):
@staticmethod
def __is_compatible():
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')
return False

View file

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