Refactored flow for main wireguard orchestration functions. Streamlined for better error handling and removal of NoneType objects

This commit is contained in:
SimplifiedPrivacy 2026-07-27 18:17:49 -04:00
parent e37d351d4c
commit 19b3925bd1
6 changed files with 186 additions and 85 deletions

View file

@ -2,6 +2,7 @@ from core.services.networking.systemwide.systemwide_wireguard import terminate_s
from core.services.networking.general_connection_tools.testing_evaluating import system_uses_wireguard_interface
from core.services.networking.general_connection_tools.connection_enable import establish_connection
# from core.services.networking.systemwide.systemwide_utils import get_firewall_setting, get_dns_setting
from core.errors.exceptions import FirewallError
from core.Errors import InvalidSubscriptionError, MissingSubscriptionError, ConnectionTerminationError, ProfileActivationError, ProfileDeactivationError, MissingLocationError, ConnectionUnprotectedError, EndpointVerificationError, ProfileStateConflictError
from core.controllers.ApplicationController import ApplicationController
@ -90,11 +91,13 @@ class ProfileController:
if profile.is_system_profile():
try:
establish_connection(profile, ignore=ignore, connection_observer=connection_observer)
except FirewallError:
raise ProfileActivationError('Firewall could not be enabled.')
except ConnectionError:
raise ProfileActivationError('The profile could not be enabled.')
if profile_observer is not None:
profile_observer.notify('enabled', profile)
if profile_observer is not None:
profile_observer.notify('enabled', profile)
@staticmethod
def disable(profile: Union[SessionProfile, SystemProfile], explicitly: bool = True, ignore: tuple[type[Exception]] = (), profile_observer: ProfileObserver = None):

View file

@ -1,4 +1,19 @@
from core.models.Result import Result, ResultError
from typing import Optional
class FirewallError(Exception):
"""There are issues with the Firewall"""
def __init__(self, result: Optional[Result]):
if result:
self.result = result
if result.message and result.message is not None:
error_msg = result.message
else:
error_msg = result.error_type
else:
error_msg = "Failed to start the firewall"
super().__init__(error_msg)
class SudoScript(Exception):
"""There are issues with the sudo scripts"""

View file

@ -112,6 +112,9 @@ def _establish_with_renegotiation(
register_wireguard_session(profile, connection_observer=connection_observer)
return establish_fn(profile, ignore=ignore, connection_observer=connection_observer)
raise ConnectionError('The connection could not be established.')
except FirewallError:
logger.error(f"Failed to turn on the firewall, passing the error up..")
raise
def __should_renegotiate(profile: Union[SessionProfile, SystemProfile]):

View file

@ -7,8 +7,9 @@ from core.services.networking.systemwide import killswitch
from core.services.networking.systemwide.dns_tools.process_dns_result import orchestrate_dns_check
from core.services.networking.systemwide.dns_tools.SearchResult import SearchResult, SearchError
from core.services.networking.systemwide.wireguard.nmcli_tools import setup_nmcli_connection, setup_ipv6_sinkhole
from core.services.networking.systemwide.wireguard.wg_firewall_dns import set_dns, revert_dns, turn_on_firewall, check_and_kill_firewall
from core.services.networking.systemwide.wireguard.wg_firewall_dns import set_dns, revert_dns, turn_on_firewall, check_and_kill_firewall, enable_firewall_with_retry
from core.errors.logger import logger
from core.errors.exceptions import FirewallError
from core.Constants import Constants
from core.Errors import ConnectionTerminationError, CommandNotFoundError
@ -80,13 +81,11 @@ def __establish_system_connection(
2) Disable any pre-existing system connections
3) Turn on WG via nmcli
4) Turn on IPv6 sinkhole
5) Turn on Firewall (if enabled)
6) Set DNS (if enabled), and test it.
7) Setup State (JSON file)
8) Confirm Firewall is on (if enabled). Raises errors if not.
9) Testing, confirming, displaying to the end user.
5) Setup State (JSON file)
6) Turn on Firewall (if enabled). Raises errors if not.
7) Set DNS (if enabled), and test it. Let's it slide even if failed.
8) Testing, confirming, displaying to the end user.
"""
check_system_prereqs()
logger.info("[CONNECT] Terming existing connections..")
terminate_system_connection(firewall_setting, dns_setting)
@ -100,7 +99,8 @@ def __establish_system_connection(
logger.info("[CONNECT] Successfully turned on WG with nmcli")
else:
logger.error(f"[SYSTEMWIDE WG] Critical issue, the initial nmcli command failed {nmcli_result.message}.")
return nmcli_result
# return nmcli_result
raise ConnectionError(f'NMCLI could not be enabled. {nmcli_result.message}')
# ============= IPv6 SINKHOLE =============
sinkhole_result = setup_ipv6_sinkhole(process_output)
@ -113,8 +113,8 @@ def __establish_system_connection(
logger.info(f"Setting State JSON..")
current_state = SystemStateController.create(
profile_id=profile.id,
firewalled=firewall_setting, # intended
dns_set=dns_setting # intended
firewalled=firewall_setting, # intended setting, not result yet
dns_set=dns_setting # intended setting, not result yet
)
# ============= PREP SETTINGS =============
@ -129,49 +129,31 @@ def __establish_system_connection(
wg_interface_name = wg_interface_result.data
# ============= FIREWALL =============
firewall_tracker = False
firewall_tracker = False # This flag exists so the JSON State can track a non-Null value, even if the firewall setting is off.
if firewall_setting:
logger.info("Firewall setting is enabled, let's turn it on..")
turn_on_result = turn_on_firewall(
logger.info("Firewall setting is enabled, attempting to enable...")
enable_firewall_with_retry(
profile_id=str(profile.id),
wg_interface_name=wg_interface_name
)
logger.info(f"Results of firewall turn on are: {turn_on_result.valid}")
if turn_on_result.valid:
firewall_tracker = True
# ============= CONFIRM ON FIREWALL =============
# if firewall_setting:
logger.info(f"Evaluating Firewall Situation after we tried to turn on it on.")
firewall_status = killswitch.status()
logger.info(f"""
The firewall setting is on.
The result of arming it, {turn_on_result.valid}
The result of the status {firewall_status}
""")
if firewall_status:
logger.info("Firewall is on. All good.")
else:
logger.error("Firewall FAILED TO ARM") # DON'T TRY TO RENEGOTIATE WITH DEAD INTERNET !!
# raise ConnectionError(f'Firewall failed to arm!! {turn_on_result.message}.') # if this triggers, it tries to renegotiate with dead internet.
firewall_tracker = True
# ============= DNS =============
if dns_setting:
set_dns_result = set_dns(str(profile.id), wg_interface_name)
logger.info(f"[CONNECT] The result of our setting the DNS is {set_dns_result.valid}")
if not set_dns_result.valid:
logger.error(f"[CONNECT] Setting the DNS DIDN'T work. Here's the reason given: {set_dns_result.message}")
# now check:
logger.error(f"[CONNECT] Setting the DNS DIDN'T work. {set_dns_result.message}")
test_dns_result = orchestrate_dns_check(wg_interface_name)
if test_dns_result.valid:
did_dns_work = True
did_dns_work = test_dns_result.valid
if did_dns_work:
logger.info("[CONNECT] All of the DNS Checks Worked")
else:
did_dns_work = False
logger.info("Skipping DNS setup, because it's off in settings")
# ============= TESTING =============
logger.info("If we made it this far, we have nmcli on, ipv6 sinkhole on, firewall on.")
try:
@ -185,16 +167,57 @@ The result of the status {firewall_status}
current_state.dns_set = did_dns_work
SystemStateController.update_or_create(current_state)
return Result(valid=True)
return Result(valid=True, data=current_state)
def evaluate_connection_result(connection_result: Result):
def evaluate_connection_result(connection_result: Result) -> None:
"""Validate the connection result. Raises if invalid."""
if not connection_result.valid:
logger.error(f"[SYSTEMWIDE WG] Critical issue, could not establish a connection: {connection_result.message}")
raise ConnectionError(f"Connection validation failed: {connection_result.message}")
def _cleanup_on_error(firewall_setting: bool, dns_setting: bool) -> None:
"""Safely terminate connection on error, swallowing cleanup errors."""
try:
terminate_system_connection(firewall_setting, dns_setting)
except ConnectionTerminationError:
pass
def _establish_connection_with_retry(
profile: SystemProfile,
firewall_setting: bool,
dns_setting: bool,
connection_observer: Optional[ConnectionObserver],
max_retries: int = 1
) -> Result:
"""Establish connection with retry logic for transient errors."""
last_error = None
for attempt in range(1, max_retries + 1):
try:
terminate_system_connection(firewall_setting, True)
except ConnectionTerminationError:
pass
connection_result = __establish_system_connection(
profile=profile,
firewall_setting=firewall_setting,
dns_setting=dns_setting,
connection_observer=connection_observer
)
evaluate_connection_result(connection_result)
return connection_result
except ConnectionError as e:
# Unrecoverable error, fail immediately
raise
except CalledProcessError as e:
# Potentially transient, retry once
last_error = e
logger.warning(f"Connection attempt {attempt}/{max_retries} failed with CalledProcessError: {e}")
if attempt == max_retries:
raise ConnectionError('The connection could not be established.') from e
def establish_system_connection(
@ -217,53 +240,74 @@ def establish_system_connection(
# ================= SETTINGS =================
firewall_setting = get_firewall_setting()
# dns_setting = True # for testing
dns_setting = get_dns_setting()
# ================= CONNECT WG ================
try:
connection_result = __establish_system_connection(
_establish_connection_with_retry(
profile=profile,
firewall_setting=firewall_setting,
dns_setting=dns_setting,
connection_observer=connection_observer
)
evaluate_connection_result(connection_result)
except ConnectionError:
# ================= RETRY ON FAILURE =================
try:
terminate_system_connection(firewall_setting, dns_setting)
except ConnectionTerminationError:
pass
raise ConnectionError('The connection could not be established.')
except CalledProcessError:
try:
terminate_system_connection(firewall_setting, dns_setting)
except ConnectionTerminationError:
pass
# trying again..
try:
connection_result = __establish_system_connection(
profile=profile,
firewall_setting=firewall_setting,
dns_setting=dns_setting,
connection_observer=connection_observer
)
evaluate_connection_result(connection_result)
except (ConnectionError, CalledProcessError):
try:
terminate_system_connection(firewall_setting, dns_setting)
except ConnectionTerminationError:
pass
raise ConnectionError('The connection could not be established.')
except (ConnectionError, CalledProcessError):
_cleanup_on_error(firewall_setting, dns_setting)
raise
terminate_tor_connection()
time.sleep(1.0)
# legacy version:
# try:
# connection_result = __establish_system_connection(
# profile=profile,
# firewall_setting=firewall_setting,
# dns_setting=dns_setting,
# connection_observer=connection_observer
# )
# evaluate_connection_result(connection_result)
# except ConnectionError:
# # ================= RETRY ON FAILURE =================
# try:
# terminate_system_connection(firewall_setting, dns_setting)
# except ConnectionTerminationError:
# pass
# raise ConnectionError('The connection could not be established.')
# except CalledProcessError:
# try:
# terminate_system_connection(firewall_setting, dns_setting)
# except ConnectionTerminationError:
# pass
# # trying again..
# try:
# connection_result = __establish_system_connection(
# profile=profile,
# firewall_setting=firewall_setting,
# dns_setting=dns_setting,
# connection_observer=connection_observer
# )
# evaluate_connection_result(connection_result)
# except (ConnectionError, CalledProcessError):
# try:
# terminate_system_connection(firewall_setting, dns_setting)
# except ConnectionTerminationError:
# pass
# raise ConnectionError('The connection could not be established.')
# def evaluate_connection_result(connection_result: Result):
# if not connection_result.valid:
# logger.error(f"[SYSTEMWIDE WG] Critical issue, could not establish a connection: {connection_result.message}")
# try:
# terminate_system_connection(firewall_setting, True)
# except ConnectionTerminationError:
# pass

View file

@ -2,6 +2,7 @@ from core.models.Result import Result, ResultError
from core.services.networking.systemwide import killswitch, dns
from core.services.networking.general_connection_tools.config_tools import extract_endpoint_ip, extract_dns
from core.errors.logger import logger
from core.errors.exceptions import FirewallError
def set_dns(profile_id: str, network_interface: str) -> Result:
do_they_even_use_systemd = dns.is_systemd_enabled()
@ -31,7 +32,7 @@ def revert_dns() -> Result:
logger.info(f"The result of the turn off is {turn_off.valid}")
if not turn_off.valid:
logger.error(f"Critical DNS Error! Could NOT turn off DNS. Error type: {turn_off.error_type} with {turn_off.message}")
return turn_off.valid
return turn_off
@ -84,6 +85,41 @@ def check_and_kill_firewall():
logger.info("We are skipping disabling the firewall, because it's already off.")
def enable_firewall_with_retry(
profile_id: str,
wg_interface_name: str,
max_retries: int = 2
) -> None:
"""Enable the firewall with automatic retry on failure."""
last_result = None
for attempt in range(1, max_retries + 1):
logger.info(f"Enabling firewall (attempt {attempt}/{max_retries})")
turn_on_result = turn_on_firewall(
profile_id=profile_id,
wg_interface_name=wg_interface_name
)
if turn_on_result is None:
logger.error("turn_on_firewall() returned None")
continue
last_result = turn_on_result
firewall_status = killswitch.status()
if turn_on_result.valid and firewall_status:
logger.info("Firewall successfully enabled")
return
logger.warning(f"Firewall enable failed on attempt {attempt}")
raise FirewallError(last_result)
# used in dead code at bottom:
# from core.services.networking.systemwide.dns_tools.process_dns_result import orchestrate_dns_check

View file

@ -59,4 +59,4 @@ requires = ["hatchling"]
build-backend = "hatchling.build"
[tool.hatch.build.targets.wheel]
packages = ["core"]-
packages = ["core"]