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

core/controllers/ProfileController.py

@@ -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,6 +91,8 @@ 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.')
 

core/errors/exceptions.py

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

core/services/networking/general_connection_tools/connection_enable.py

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

core/services/networking/systemwide/systemwide_wireguard.py

@@ -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)
+    5) Setup State (JSON file)
-    6) Set DNS (if enabled), and test it.
+    6) Turn on Firewall (if enabled). Raises errors if not.
-    7) Setup State (JSON file)
+    7) Set DNS (if enabled), and test it. Let's it slide even if failed.
-    8) Confirm Firewall is on (if enabled).  Raises errors if not.
+    8) Testing, confirming, displaying to the end user.
-    9) 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
+        firewalled=firewall_setting, # intended setting, not result yet
-        dns_set=dns_setting # intended
+        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..")
+        logger.info("Firewall setting is enabled, attempting to enable...")
-        turn_on_result = turn_on_firewall(
+        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.
-
     # ============= 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}")
+            logger.error(f"[CONNECT] Setting the DNS DIDN'T work. {set_dns_result.message}")
         
-        # now check:
         test_dns_result = orchestrate_dns_check(wg_interface_name)
-        if test_dns_result.valid:
+        did_dns_work = test_dns_result.valid
-            did_dns_work = True
+        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,18 +167,59 @@ 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, True)
+        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:
+            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(
     profile: SystemProfile,
     ignore: tuple[type[Exception]] = (),
@@ -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):
-
+        _cleanup_on_error(firewall_setting, dns_setting)
-            try:
+        raise
-                terminate_system_connection(firewall_setting, dns_setting)
-            except ConnectionTerminationError:
-                pass
-
-            raise ConnectionError('The connection could not be established.')
 
     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

core/services/networking/systemwide/wireguard/wg_firewall_dns.py

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

pyproject.toml

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