From f53661fc60c6d1671181967adefbf85cad795089 Mon Sep 17 00:00:00 2001 From: SimplifiedPrivacy Date: Sun, 25 Oct 2026 20:55:19 -0400 Subject: [PATCH] Improved Systemwide Wireguard Error Flow, to cache each stage of progress to the end of the loop, and prompt the user on each choice. --- core/models/ConnectionRecovery.py | 66 +++++ core/models/Result.py | 1 + .../keys_and_verifications/wireguard_keys.py | 112 +------- .../networking/general_connection_tools/ip.py | 247 +---------------- .../networking/mtu/mtu_orchestrator.py | 10 +- .../networking/ports/port_orchestrator.py | 149 ++++++++++ .../systemwide/systemwide_wireguard.py | 257 +++++++++++++----- core/services/payment_phase/payment_state.py | 203 +++++++------- core/utils/singleton.py | 25 -- core/utils/singleton_state.py | 161 +++++++++++ 10 files changed, 684 insertions(+), 547 deletions(-) create mode 100644 core/models/ConnectionRecovery.py create mode 100644 core/services/networking/ports/port_orchestrator.py delete mode 100644 core/utils/singleton.py create mode 100644 core/utils/singleton_state.py diff --git a/core/models/ConnectionRecovery.py b/core/models/ConnectionRecovery.py new file mode 100644 index 0000000..90eb490 --- /dev/null +++ b/core/models/ConnectionRecovery.py @@ -0,0 +1,66 @@ +from core.utils.singleton_state import singleton, StatePersistenceMixin, wipe_model +from core.Constants import Constants +from core.errors.logger import logger + +from dataclasses import dataclass +import os +from functools import partial + +@singleton +@dataclass +class ConnectionRecovery(StatePersistenceMixin): + PAYLOAD_FIELDS = {"port_changed", "port_changed_to", "port_change_worked", "mtu_tested", "mtu_changed", "wg_key_changed", "location_tested", "location_up"} + KEEP_FIELDS = ["parent_folder", "instance_id", "folder_path", "_data_file"] + + # metadata on it's own instance: + parent_folder = None + instance_id: int = None + folder_path: str = None + _data_file: str = None + + # actual payload data: + + # LOCATION + location_tested: bool = False + location_up: bool = False + + # PORT + asked_to_change_port: bool = False + port_changed: bool = False + port_changed_to: int|None = None + port_change_worked: bool = False + + # MTU + mtu_tested: bool = False + mtu_changed: bool = False + asked_to_change_mtu: bool = False + + # WG KEY + asked_to_change_wg_key: bool = False + wg_key_confirmed: bool = False + wg_key_changed: bool = False + + # Overall state + attempts: int = 0 + bypass_attempts: bool = False + reached_end_of_loop: bool = False + + stashed_partial: partial = None + + def __post_init__(self): + # Localized mode: each ID gets its own folder + self.parent_folder = Constants.HV_PROFILE_CONFIG_HOME + self.folder_path = os.path.join(self.parent_folder, str(self.instance_id)) + os.makedirs(self.folder_path, exist_ok=True) + self._data_file = os.path.join(self.folder_path, "recovery.json") + + logger.info(f"full path of self._data_file {self._data_file}") + + self.load() + + def wipe(self): + """Wipe all data and start over.""" + wipe_model(model=self, KEEP_FIELDS=self.KEEP_FIELDS) + wiped = self.wipe_data_on_disk() + logger.info(f"Wiped: {wiped}") + diff --git a/core/models/Result.py b/core/models/Result.py index dcb2972..f9920d4 100644 --- a/core/models/Result.py +++ b/core/models/Result.py @@ -58,6 +58,7 @@ class ResultError(Enum): NO_INTERNET = "no_internet" PORT = "port" RENEGOTIATION = "renegotiation" + ATTEMPTS = "attempts" NAVIGATE_TO = "navigate_to" UNKNOWN = "unknown" diff --git a/core/services/keys_and_verifications/wireguard_keys.py b/core/services/keys_and_verifications/wireguard_keys.py index a2647f2..b823211 100644 --- a/core/services/keys_and_verifications/wireguard_keys.py +++ b/core/services/keys_and_verifications/wireguard_keys.py @@ -1,3 +1,4 @@ +from core.models.ConnectionRecovery import ConnectionRecovery from core.controllers.ConfigurationController import ConfigurationController from core.infrastructure.global_state import ProfileStateRepository from core.utils.get_raw_string import get_string @@ -29,8 +30,6 @@ from functools import partial MIN_MTU = 1280 -PORT_CHOICES = [48271, 52428, 51820] - def register_wireguard_session( profile: Profile, connection_observer: Optional[ConnectionObserver] = None @@ -189,112 +188,3 @@ def expand_newlines(text): """ return "\n\n".join([line for line in text.splitlines() if line.strip()]) - -def get_port(config: str) -> int: - """Get the wireguard port if you have a config already.""" - # This regex searches for the "Endpoint" key, ignores optional whitespace around the '=', - # skips the IP/hostname part until the colon, and captures the digits following it. - match = re.search(r"Endpoint\s*=\s*[^:]+:(\d+)", config, re.IGNORECASE) - if match: - return int(match.group(1)) - return None - -def change_port(new_port: int, config: str) -> str: - # This regex captures everything from the start of the line up to the colon - # in the first group, then matches the existing digits of the port. - pattern = r"(Endpoint\s*=\s*[^:]+):(\d+)" - - # We use a lambda function as the replacement to insert the new_port variable - # while preserving the captured IP/hostname in group 1 (\1). - return re.sub(pattern, lambda m: f"{m.group(1)}:{new_port}", config, flags=re.IGNORECASE) - - -def get_wg_port(profile: Profile) -> Result: - """Get the Wireguard port with a given profile.""" - config_path = profile.get_wireguard_configuration_path() - if not config_path: - return Result(valid=False, error_type=ResultError.INVALID_INPUT, message="Can't get the path of the config") - - config_data = get_string(config_path) - if not config_data: - return Result(valid=False, error_type=ResultError.MISSING_DATA, message="Missing the config data, or could not read it.") - - port = get_port(config_data) - if port: - return Result(valid=True, data=port) - else: - return Result(valid=False, error_type=ResultError.INVALID_INPUT, message="Config is not in the correct format to be read.") - - -def get_random_different_port(current_port: int) -> int|None: - """Get a different port (that's available) from the given choices.""" - while True: - new_port = random.choice(PORT_CHOICES) - if str(new_port) != str(current_port): - return new_port - - -def change_wg_port(profile: Profile = None, desired_port: Optional[int] = None) -> Result: - """Change a profile's Wireguard port.""" - - logger.info(f"desired_port is {desired_port}") - config_path = profile.get_wireguard_configuration_path() - config_data = get_string(config_path) - current_port = get_port(config_data) - logger.info(f"current_port is {current_port}") - - if desired_port and desired_port != current_port and desired_port in PORT_CHOICES: - change_to = desired_port - else: - change_to = get_random_different_port(current_port) - logger.info(f"Changing the port to {change_to}") - - new_config = change_port(new_port=change_to, config=config_data) - - # sudo prompts: - try: - profile.attach_wireguard_configuration(new_config) - logger.info("changed the port. updating state..") - - # UPDATE STATE: - state_repo = ProfileStateRepository() - state_repo.global_state.changed_port = True - - return Result(valid=True, data=change_to) - except ProfileModificationError: - return Result(valid=False, error_type=ResultError.PERMISSON, message="Failed to swap the port due to permissions") - except CommandNotFoundError: - return Result(valid=False, error_type=ResultError.MISSING_SOFTWARE, message="Missing the command") - - -def port_was_changed_already_today() -> bool: - state_repo = ProfileStateRepository() - return state_repo.global_state.changed_port - - -def offer_port_change(profile: Profile) -> Result: - """Offer the user the option to change the port""" - port_result = get_wg_port(profile) - - if not port_result.valid: - return port_result - - current_port = port_result.data - desired_port = get_random_different_port(current_port) - - func_choice = partial(change_wg_port, profile=profile, desired_port=desired_port) - return Result(valid=False, function_choice=func_choice, error_type=ResultError.PORT, message=f"Failed to connect. Do you want to try to change the port to {desired_port}?") - - -def should_we_try_to_change_port(): - """Returns a boolean if we should change the port""" - if not ConfigurationController.do_we_allow_port_prompts(): - logger.info("We do not allow port changes") - return False - - state_repo = ProfileStateRepository() - if state_repo.global_state.changed_port: - return False - else: - return True - diff --git a/core/services/networking/general_connection_tools/ip.py b/core/services/networking/general_connection_tools/ip.py index 7eebc6a..3a6a995 100644 --- a/core/services/networking/general_connection_tools/ip.py +++ b/core/services/networking/general_connection_tools/ip.py @@ -24,6 +24,10 @@ import httpx_socks from httpx_socks import ProxyType import subprocess from typing import Optional +import socket +import multiprocessing +from threading import Thread +from queue import Queue, Empty API_ENDPOINT_FORMATS = { f"{Constants.SP_BASE_URL}/v2/get_ip": 'ip', @@ -38,16 +42,8 @@ API_ENDPOINT_FORMATS = { MAX_ATTEMPTS = 3 - - -import time -import socket -import multiprocessing -from threading import Thread -from queue import Queue, Empty - def is_dns_problem_worker(url: str, queue: multiprocessing.Queue): - """Worker function for a separate PROCESS""" + """Checks if DNS is the problem. (It's a Worker function for a separate process)""" domain = extract_domain(url) try: socket.gethostbyname(domain) @@ -60,7 +56,7 @@ def test(connection_observer: Optional[ConnectionObserver], port: Optional[int] start_time = time.time() total_budget = 5.0 - # 1. NUCLEAR DNS PRE-FLIGHT (Using Processes, not Threads) + # 1. DNS PRE-FLIGHT (Using Processes, not Threads) if port is None: dns_queue = multiprocessing.Queue() processes = [] @@ -71,7 +67,7 @@ def test(connection_observer: Optional[ConnectionObserver], port: Optional[int] processes.append(p) any_resolved = False - dns_timeout = 2 + dns_timeout = 3 try: # Wait for the first successful resolution @@ -85,7 +81,6 @@ def test(connection_observer: Optional[ConnectionObserver], port: Optional[int] except Empty: continue finally: - # THE NUCLEAR STRIKE: # We kill every single process immediately. # The OS kernel wipes them from memory regardless of where they are stuck. for p in processes: @@ -144,232 +139,6 @@ def test(connection_observer: Optional[ConnectionObserver], port: Optional[int] return Result(valid=False, error_type=ResultError.CONNECTION, message="Could not connect.") - - - - - - - -# import time -# import socket -# from concurrent.futures import ThreadPoolExecutor, wait, FIRST_COMPLETED - -# def is_dns_problem(url: str) -> bool: -# """Regular DNS check for domains""" -# domain = extract_domain(url) -# try: -# socket.gethostbyname(domain) -# return False -# except (socket.gaierror, Exception): -# return True - -# def test(connection_observer: Optional[ConnectionObserver], port: Optional[int] = None) -> Result: -# endpoint_list = list(API_ENDPOINT_FORMATS.keys())[:MAX_ATTEMPTS] -# start_time = time.time() -# total_budget = 4.0 - -# # 1. AGGRESSIVE DNS PRE-FLIGHT -# if port is None: -# # IMPORTANT: Do NOT use 'with' here. 'with' calls shutdown(wait=True), -# # which causes the hang you are seeing. -# dns_executor = ThreadPoolExecutor(max_workers=len(endpoint_list)) -# dns_futures = {dns_executor.submit(is_dns_problem, url): url for url in endpoint_list} - -# dns_done, _ = wait(dns_futures, timeout=1.2, return_when=FIRST_COMPLETED) - -# any_resolved = False -# for future in dns_done: -# if future.result() is False: -# any_resolved = True -# break - -# if not any_resolved: -# logger.error("Aggressive DNS check failed: No internet or DNS unreachable.") -# # Return immediately. The dns_executor threads will stay alive -# # in the background until the OS times them out, but the -# # main thread is now free. -# return Result(valid=False, error_type=ResultError.CONNECTION, message="No internet connection.") - -# # 2. HTTP REQUEST PHASE -# transport = None -# if port: -# transport = httpx_socks.ProxyTransport( -# proxy_type=ProxyType.SOCKS5, -# proxy_host="127.0.0.1", -# proxy_port=port, -# ) - -# with httpx.Client(transport=transport, http2=False, timeout=3.0) as client: -# def attempt_request(url: str, index: int): -# try: -# raw_response = client.get(url) -# filtered_response = classify_response(raw_response) -# if filtered_response.valid: -# ip_address = extract_ip_from_response(response=filtered_response.data, api_url=url) -# if ip_address: -# return Result(valid=True, data=ip_address) -# except Exception: -# pass -# return Result(valid=False) - -# # IMPORTANT: Do NOT use 'with' here either. -# http_executor = ThreadPoolExecutor(max_workers=len(endpoint_list)) -# http_futures = {http_executor.submit(attempt_request, url, i): url for i, url in enumerate(endpoint_list)} - -# while http_futures: -# elapsed = time.time() - start_time -# remaining = total_budget - elapsed - -# if remaining <= 0: -# break - -# done, _ = wait(http_futures, timeout=remaining, return_when=FIRST_COMPLETED) - -# if not done: -# break - -# for future in done: -# result = future.result() -# if result.valid: -# return result -# http_futures.remove(future) - -# logger.error(f"We went through {len(endpoint_list)} attempts and could not get a reply.") -# return Result(valid=False, error_type=ResultError.CONNECTION, message="Could not connect.") - - - - - - - - - -# async - - - -# async def _test(connection_observer: Optional[ConnectionObserver], port: Optional[int] = None) -> Result: -# endpoint_list = list(API_ENDPOINT_FORMATS.keys())[:MAX_ATTEMPTS] - -# # Using 'async with' ensures the client and its connection pool are -# # closed before the event loop starts shutting down. -# async with httpx.AsyncClient( -# transport=httpx_socks.AsyncProxyTransport( -# proxy_type=ProxyType.SOCKS5, -# proxy_host="127.0.0.1", -# proxy_port=port, -# ) if port else None, -# http2=True, -# timeout=4 -# ) as client: - -# async def attempt_request(url: str, index: int): -# logger.info(f"Attempt {index + 1}: Trying {url}") -# try: -# raw_response = await client.get(url) -# filtered_response = classify_response(raw_response) - -# if filtered_response.valid: -# ip_address = extract_ip_from_response( -# response=filtered_response.data, -# api_url=url -# ) -# if ip_address is not None: -# return Result(valid=True, data=ip_address) - -# logger.error(f"Skipping {url}") -# except Exception as e: -# logger.error(f"Skipping {url}. Invalid API reply or failed connection {str(e)}") - -# return Result(valid=False) - -# # Create tasks for concurrent execution -# tasks = [asyncio.create_task(attempt_request(url, i)) for i, url in enumerate(endpoint_list)] - -# try: -# for future in asyncio.as_completed(tasks): -# result = await future -# if result.valid: -# # 1. Signal all other pending tasks to stop -# for t in tasks: -# if not t.done(): -# t.cancel() - -# # 2. CRITICAL: Wait for the cancelled tasks to actually finish. -# # return_exceptions=True prevents CancelledError from being raised here. -# await asyncio.gather(*tasks, return_exceptions=True) -# return result -# except Exception as e: -# logger.error(f"Unexpected error during request processing: {e}") - -# # This part is reached only if no task returned a valid Result -# logger.error(f"We went through {len(endpoint_list)} attempts and could not get a reply.") -# return Result(valid=False, error_type=ResultError.CONNECTION, message="Could not connect.") - - -# def test(connection_observer: Optional[ConnectionObserver], port: Optional[int] = None) -> Result: -# if connection_observer: -# connection_observer.notify("connecting", "Testing IP..") -# return asyncio.run(_test(connection_observer, port)) - - -# async def _test(connection_observer: Optional[ConnectionObserver], port: Optional[int] = None) -> Result: -# endpoint_list = list(API_ENDPOINT_FORMATS.keys())[:MAX_ATTEMPTS] - -# if port: -# transport = httpx_socks.AsyncProxyTransport( -# proxy_type=ProxyType.SOCKS5, -# proxy_host="127.0.0.1", -# proxy_port=port, -# ) -# client = httpx.AsyncClient(transport=transport, http2=True, timeout=4) -# else: -# client = httpx.AsyncClient(http2=True, timeout=4) - -# async def attempt_request(url: str, index: int): -# logger.info(f"Attempt {index + 1}: Trying {url}") - -# try: -# raw_response = await client.get(url) -# filtered_response = classify_response(raw_response) - -# if filtered_response.valid: -# ip_address = extract_ip_from_response( -# response=filtered_response.data, -# api_url=url -# ) -# if ip_address is not None: -# return Result(valid=True, data=ip_address) - -# logger.error(f"Skipping {url}") -# except Exception as e: -# logger.error(f"Skipping {url}. Invalid API reply or failed connection {str(e)}") - -# return Result(valid=False) - -# try: -# # Create tasks for all endpoints to run them concurrently -# tasks = [asyncio.create_task(attempt_request(url, i)) for i, url in enumerate(endpoint_list)] - -# # Process tasks as they complete -# for future in asyncio.as_completed(tasks): -# result = await future -# if result.valid: -# # Cancel all other pending tasks immediately -# for t in tasks: -# if not t.done(): -# t.cancel() -# return result - -# finally: -# await client.aclose() - -# logger.error(f"We went through {len(endpoint_list)} attempts and could not get a reply.") -# return Result(valid=False, error_type=ResultError.CONNECTION, message="Could not connect.") - - def extract_ip_from_response(response: dict, api_url: str) -> str | bool: """ Purpose: @@ -398,6 +167,7 @@ def extract_ip_from_response(response: dict, api_url: str) -> str | bool: def ping_ip(ip: str) -> Result: + """Ping an IP using the pre-existing terminal""" command = ["ping", "-c", "1", "-W", "2", ip] human_readable_goal = f"Pinging {ip}" Notifications.emit(EventType.STATUS_UPDATE, human_readable_goal) @@ -406,6 +176,7 @@ def ping_ip(ip: str) -> Result: def ping_ip_worked(ip: str) -> bool: + """Test ping in a new subprocess""" try: # -c 1: send only 1 packet # -W 2: wait 2 seconds for a response diff --git a/core/services/networking/mtu/mtu_orchestrator.py b/core/services/networking/mtu/mtu_orchestrator.py index 835456d..b7c35a2 100644 --- a/core/services/networking/mtu/mtu_orchestrator.py +++ b/core/services/networking/mtu/mtu_orchestrator.py @@ -1,3 +1,4 @@ +from core.models.ConnectionRecovery import ConnectionRecovery from core.services.keys_and_verifications.wireguard_keys import get_mtu, insert_mtu_into_config_on_disk from core.services.networking.mtu.mtu_tools import loop_through_testing_mtu, mtu_is_fine, MIN_MTU, MAX_MTU from core.observers.Notifications import Notifications, EventType @@ -27,7 +28,10 @@ def _changed_mtu(profile: SystemProfile, current_mtu: Optional[int]) -> bool: Notifications.emit(EventType.STATUS_UPDATE, notification) return insert_mtu_into_config_on_disk(profile=profile, mtu=optimal_mtu) -def should_we_even_test_mtu(profile_id: int): +def should_we_even_test_mtu(profile_id: int, recovery_stats: ConnectionRecovery) -> bool: + if recovery_stats.mtu_tested: + return False + # Evaluate MTU WITH connection on current_mtu = get_mtu(profile_id=profile_id) @@ -41,7 +45,7 @@ def should_we_even_test_mtu(profile_id: int): logger.info("MTU isn't at the min, so we can potentially improve") return True, current_mtu -def was_mtu_changed(profile: Profile, current_mtu: int): +def was_mtu_changed(profile: Profile, current_mtu: int, recovery_stats: ConnectionRecovery) -> bool: notification = "Testing MTU..." Notifications.emit(EventType.STATUS_UPDATE, notification) logger.info(notification) @@ -56,6 +60,8 @@ def was_mtu_changed(profile: Profile, current_mtu: int): if _changed_mtu(profile=profile, current_mtu=current_mtu): logger.info("MTU changed. We should retry the connection") + recovery_stats.mtu_changed = True + recovery_stats.save() return True else: diff --git a/core/services/networking/ports/port_orchestrator.py b/core/services/networking/ports/port_orchestrator.py new file mode 100644 index 0000000..86ae98a --- /dev/null +++ b/core/services/networking/ports/port_orchestrator.py @@ -0,0 +1,149 @@ +from core.models.ConnectionRecovery import ConnectionRecovery +from core.infrastructure.global_state import ProfileStateRepository +from core.controllers.ConfigurationController import ConfigurationController +from core.utils.get_raw_string import get_string +from core.Errors import ProfileModificationError, CommandNotFoundError +from core.models.BaseProfile import BaseProfile as Profile +from core.models.Result import Result, ResultError +from core.errors.logger import logger + +import re +import random +# from functools import partial +from typing import Optional + +PORT_CHOICES = [48271, 52428, 51820] + + +def get_port(config: str) -> int: + """Get the wireguard port if you have a config already.""" + # This regex searches for the "Endpoint" key, ignores optional whitespace around the '=', + # skips the IP/hostname part until the colon, and captures the digits following it. + match = re.search(r"Endpoint\s*=\s*[^:]+:(\d+)", config, re.IGNORECASE) + if match: + return int(match.group(1)) + return None + +def change_port(new_port: int, config: str) -> str: + # This regex captures everything from the start of the line up to the colon + # in the first group, then matches the existing digits of the port. + pattern = r"(Endpoint\s*=\s*[^:]+):(\d+)" + + # We use a lambda function as the replacement to insert the new_port variable + # while preserving the captured IP/hostname in group 1 (\1). + return re.sub(pattern, lambda m: f"{m.group(1)}:{new_port}", config, flags=re.IGNORECASE) + + +def get_wg_port(profile: Profile) -> Result: + """Get the Wireguard port with a given profile.""" + config_path = profile.get_wireguard_configuration_path() + if not config_path: + return Result(valid=False, error_type=ResultError.INVALID_INPUT, message="Can't get the path of the config") + + config_data = get_string(config_path) + if not config_data: + return Result(valid=False, error_type=ResultError.MISSING_DATA, message="Missing the config data, or could not read it.") + + port = get_port(config_data) + if port: + return Result(valid=True, data=port) + else: + return Result(valid=False, error_type=ResultError.INVALID_INPUT, message="Config is not in the correct format to be read.") + + +def get_random_different_port(current_port: int) -> int|None: + """Get a different port (that's available) from the given choices.""" + while True: + new_port = random.choice(PORT_CHOICES) + if str(new_port) != str(current_port): + return new_port + + +def change_wg_port(profile: Profile, desired_port: Optional[int] = None) -> Result: + """Change a profile's Wireguard port.""" + + logger.info(f"Starting to change the port for profile {profile.id}") + + # pre-existing data: + recovery_stats = ConnectionRecovery() + + # get config path: + config_path = profile.get_wireguard_configuration_path() + + # get config data: + config_data = get_string(config_path) + + # assess current port on that data: + current_port = get_port(config_data) + logger.info(f"current_port is {current_port}") + + # changing to what? + if not desired_port: + desired_port = recovery_stats.port_changed_to + else: + desired_port = get_random_different_port(current_port) + + # sanity filter: + if desired_port == current_port or desired_port not in PORT_CHOICES: + desired_port = get_random_different_port(current_port) + + logger.info(f"Changing the port to {desired_port}") + new_config = change_port(new_port=desired_port, config=config_data) + + # sudo prompts: + try: + profile.attach_wireguard_configuration(new_config) + logger.info("changed the port. updating state..") + + # UPDATE STATE: + recovery_stats.port_changed = True + recovery_stats.save() + + state_repo = ProfileStateRepository() + state_repo.global_state.changed_port = True + + return Result(valid=True, data=desired_port) + except ProfileModificationError: + return Result(valid=False, error_type=ResultError.PERMISSION, message="Failed to swap the port due to permissions") + except CommandNotFoundError: + return Result(valid=False, error_type=ResultError.MISSING_SOFTWARE, message="Missing the command") + + +def port_was_changed_already_today() -> bool: + state_repo = ProfileStateRepository() + return state_repo.global_state.changed_port + + +def should_we_try_to_change_port(recovery_stats: ConnectionRecovery): + """Returns a boolean if we should change the port""" + + # already changed: + if recovery_stats.port_changed or recovery_stats.asked_to_change_port: + return False + + if not ConfigurationController.do_we_allow_port_prompts(): + logger.info("We do not allow port changes") + return False + + state_repo = ProfileStateRepository() + if state_repo.global_state.changed_port: + return False + else: + return True + + +# def offer_port_change(profile: Profile, recovery_stats: ConnectionRecovery) -> Result: +# """Offer the user the option to change the port""" +# port_result = get_wg_port(profile) + +# if not port_result.valid: +# return port_result + +# current_port = port_result.data +# desired_port = get_random_different_port(current_port) +# recovery_stats.port_changed_to = desired_port +# recovery_stats.save() + +# # func_choice = partial(change_wg_port, profile=profile, desired_port=desired_port) +# return Result(valid=False, error_type=ResultError.PORT, message=f"Failed to connect. Do you want to try to change the port to {desired_port}?") + diff --git a/core/services/networking/systemwide/systemwide_wireguard.py b/core/services/networking/systemwide/systemwide_wireguard.py index 7ac9394..55fddb8 100644 --- a/core/services/networking/systemwide/systemwide_wireguard.py +++ b/core/services/networking/systemwide/systemwide_wireguard.py @@ -1,6 +1,8 @@ +from core.models.ConnectionRecovery import ConnectionRecovery from core.services.keys_and_verifications.wireguard_keys import register_wireguard_session -from core.services.keys_and_verifications.wireguard_keys import offer_port_change, should_we_try_to_change_port #, port_was_changed_already_today +from core.services.networking.ports.port_orchestrator import get_wg_port, should_we_try_to_change_port, get_random_different_port, change_wg_port from core.services.networking.mtu.mtu_orchestrator import should_we_even_test_mtu, was_mtu_changed +from core.models.BaseProfile import BaseProfile as Profile from core.services.networking.general_connection_tools.location_evaluation import check_if_location_is_up from core.services.networking.api_requests.ApiResponseModel import BackoffStrategy @@ -25,7 +27,7 @@ from core.utils.run_commands import run_generic_command from core.observers.TicketObserver import TicketObserver # from core.Constants import Constants -from core.Errors import ConnectionTerminationError, CommandNotFoundError +from core.Errors import CommandNotFoundError # ConnectionTerminationError from core.controllers.ConfigurationController import ConfigurationController from core.controllers.SystemStateController import SystemStateController # from core.models.BaseProfile import ProfileType @@ -43,6 +45,8 @@ import time # import asyncio from functools import partial +MAX_ATTEMPTS = 4 + def terminate_system_connection( firewall_setting: Optional[bool] = None, dns_setting: Optional[bool] = None @@ -105,6 +109,129 @@ def terminate_system_connection( return Result(valid=True) +def renegotiate_and_reconnect( + profile: SystemProfile, + ignore: tuple[type[Exception]] = (), + connection_observer: Optional[ConnectionObserver] = None, + ticket_observer: Optional[TicketObserver] = None +) -> Result: + """Register a new wireguard session with a new key, and connect.""" + + notification = "Renegotating WG Key.." + Notifications.emit(EventType.STATUS_UPDATE, notification) + logger.info(notification) + + register_wireguard_session(profile=profile, connection_observer=connection_observer) + logger.info("Registered new session") + + recovery_stats = ConnectionRecovery() + recovery_stats.wg_key_changed = True + + notification = "Reconnecting with new WG Key.." + Notifications.emit(EventType.STATUS_UPDATE, notification) + logger.info(notification) + + return establish_system_connection( + profile=profile, + ignore=ignore, + connection_observer=connection_observer, + ticket_observer=ticket_observer + ) + + +def wipe_attempts( + recovery_stats: ConnectionRecovery, + profile: SystemProfile, + ignore: tuple[type[Exception]] = (), + connection_observer: Optional[ConnectionObserver] = None, + ticket_observer: Optional[TicketObserver] = None +)-> Result: + """Wipe pre-existing recovery attempts and solutions. And begin fresh.""" + recovery_stats.wipe() + + notification = "Starting Recovery Fresh" + Notifications.emit(EventType.STATUS_UPDATE, notification) + logger.info(notification) + + return establish_system_connection( + profile=profile, + ignore=ignore, + connection_observer=connection_observer, + ticket_observer=ticket_observer + ) + + +def bypass_existing_attempts( + recovery_stats: ConnectionRecovery, + profile: SystemProfile, + ignore: tuple[type[Exception]] = (), + connection_observer: Optional[ConnectionObserver] = None, + ticket_observer: Optional[TicketObserver] = None +)-> Result: + """Stop prompting the user for solutions (ports, mtu, wg-renegotiation)""" + + notification = "Skipping Recovery Strategies" + Notifications.emit(EventType.STATUS_UPDATE, notification) + logger.info(notification) + + recovery_stats.bypass_attempts = True + recovery_stats.reached_end_of_loop = False + recovery_stats.save() + + return establish_system_connection( + profile=profile, + ignore=ignore, + connection_observer=connection_observer, + ticket_observer=ticket_observer + ) + + +def change_port_then_reconnect( + profile: SystemProfile, + ignore: tuple[type[Exception]] = (), + connection_observer: Optional[ConnectionObserver] = None, + ticket_observer: Optional[TicketObserver] = None +) -> Result: + + notification = "Changing the port.." + Notifications.emit(EventType.STATUS_UPDATE, notification) + logger.info(notification) + + changed = change_wg_port(profile) + + if not changed.valid: + return changed + + notification = "Reconnecting with New Port.." + Notifications.emit(EventType.STATUS_UPDATE, notification) + logger.info(notification) + + return establish_system_connection( + profile=profile, + ignore=ignore, + connection_observer=connection_observer, + ticket_observer=ticket_observer + ) + +def get_port_change_message(profile: Profile, recovery_stats: ConnectionRecovery) -> str: + """Offer the user the option to change the port""" + port_result = get_wg_port(profile) + + if not port_result.valid: + return port_result + + current_port = port_result.data + desired_port = get_random_different_port(current_port) + recovery_stats.port_changed_to = desired_port + recovery_stats.asked_to_change_port = True + recovery_stats.save() + + notification = "Change Port?" + Notifications.emit(EventType.STATUS_UPDATE, notification) + + message = f"Failed to connect. Do you want to try to change the port to {desired_port}?" + return message + def evaluate_location_and_internet(profile: SystemProfile, connection_observer: ConnectionObserver) -> Result: """Checks if location is up and returns the backoff strategy.""" @@ -119,10 +246,6 @@ def evaluate_location_and_internet(profile: SystemProfile, connection_observer: return location_result -# def _cleanup_on_error(firewall_setting: bool, dns_setting: bool) -> None: -# """Safely terminate connection on error, swallowing cleanup errors.""" -# terminate_system_connection(firewall_setting, dns_setting) - def _finish_success(final_result: Result, profile_id: int) -> Result: terminate_tor_connection() @@ -130,11 +253,25 @@ def _finish_success(final_result: Result, profile_id: int) -> Result: return final_result +def _conclude_loop_failure(): + """Finish the loop with a generic error message""" + + logger.info("Saving the loop has concluded.") + recovery_stats = ConnectionRecovery() + recovery_stats.reached_end_of_loop = True + recovery_stats.save() + + notification = "Connection failed" + Notifications.emit(EventType.STATUS_UPDATE, notification) + error_msg = "Connection failed, but the location is up. We recommend you try a different protocol." + return Result(valid=False, error_type=ResultError.UNKNOWN, message=error_msg) + + def __establish_system_connection( profile: SystemProfile, firewall_setting: bool = False, dns_setting: bool = False, - connection_observer: Optional[ConnectionObserver] = None): + connection_observer: Optional[ConnectionObserver] = None) -> Result: """ Purpose: Orchestrates the full flow of the connection. @@ -216,8 +353,6 @@ def __establish_system_connection( # ============= TESTING ============= logger.info("If we made it this far, we have nmcli on, ipv6 sinkhole on, firewall on.") - # if connection_observer: - # connection_observer.notify("connecting", "Testing Connection Health..") notification = "Testing IP address.." Notifications.emit(EventType.STATUS_UPDATE, notification) @@ -232,14 +367,6 @@ def __establish_system_connection( logger.error("Returning that the test failed..") return test_result - # ============= UPDATE STATE ============= - # if dns_setting and not did_dns_work: - # current_state.firewalled = firewall_tracker - # current_state.dns_set = did_dns_work - # SystemStateController.update_or_create(current_state) - # notification = "DNS Set Did Not Work." - # Notifications.emit(EventType.STATUS_UPDATE, notification) - # ============= SETUP STATE ============= logger.info("Setting State JSON..") current_state = SystemStateController.create( @@ -250,7 +377,6 @@ def __establish_system_connection( ) return Result(valid=True, data=current_state) -import time def establish_system_connection( profile: SystemProfile, @@ -259,7 +385,7 @@ def establish_system_connection( ticket_observer: Optional[TicketObserver] = None # Passed in only for assassin's link to tickets in ConnectionController.establish_session_connection ) -> Result|None: """ - Establish connection with a maximum of 2 retries. + Establish connection. Steps) 1) Check Prereqs @@ -279,10 +405,19 @@ def establish_system_connection( # ================= SETTINGS ================= firewall_setting = ConfigurationController.get_firewall_setting() dns_setting = ConfigurationController.get_dns_setting() + recovery_stats = ConnectionRecovery(instance_id=profile.id) + + if recovery_stats.reached_end_of_loop: + error_msg = "Your previous attempt to enable this profile failed. Do you want to start fresh with your recovery strategies (recommended. if the connection fails then it will check & reask about port, renegotation, and mtu)" + yes_choice = partial(wipe_attempts, recovery_stats=recovery_stats, profile=profile, ignore=ignore, connection_observer=connection_observer, ticket_observer=ticket_observer) + no_option = partial(bypass_existing_attempts, recovery_stats=recovery_stats, profile=profile, ignore=ignore, connection_observer=connection_observer, ticket_observer=ticket_observer) + return Result(valid=False, function_choice=yes_choice, no_option=no_option, error_type=ResultError.ATTEMPTS, message=error_msg) # ================= CONNECT WG ================ - # We allow 3 total attempts (Initial + 2 retries) - for attempt in range(1, 4): + while not recovery_stats.reached_end_of_loop: + recovery_stats.attempts += 1 + current_try = recovery_stats.attempts + # ================================= # PRE-EXISTING # ================================= @@ -294,13 +429,7 @@ def establish_system_connection( logger.error(f"Failed to terminate pre-existing {terminated.message}") return terminated - # func_choice = partial(register_wireguard_session, profile=profile, connection_observer=connection_observer) - # return Result(valid=False, function_choice=func_choice, error_type=ResultError.RENEGOTIATION, message="Would you like to try Wireguard key renegotiation? (Note: this invalidates other profiles with the same subscription.)") - - print("sleeping") - time.sleep(5) - - notification = f"Connecting attempt {attempt}.. " + notification = f"Connecting attempt {current_try}.. " Notifications.emit(EventType.STATUS_UPDATE, notification) # ================================= @@ -314,19 +443,17 @@ def establish_system_connection( connection_observer=connection_observer ) - print("sleeping") - time.sleep(5) - # ================================= # EVALUATE # ================================= if connection_result.valid: + recovery_stats.wipe() return _finish_success(final_result=connection_result, profile_id=profile.id) # ================================= # ATTEMPT ONE: LOCATION UP? # ================================= - if attempt == 1: + if current_try == 1 and not recovery_stats.location_tested: # First failure: Evaluate Location and Internet without connection, logger.info("Terminating system connection to check location") terminate_system_connection(firewall_setting, dns_setting) @@ -335,9 +462,11 @@ def establish_system_connection( profile=profile, connection_observer=connection_observer ) + recovery_stats.location_tested = True #============== NO RETRY #============== if location_result.backoff_strategy == BackoffStrategy.NO_RETRY: + recovery_stats.location_up = False logger.error(f"Location check failed: {location_result.message}") bottom_left = "Node Down" Notifications.emit(EventType.STATUS_UPDATE, bottom_left) @@ -347,6 +476,7 @@ def establish_system_connection( #============== NO SEEMS FINE #============== notification = "Location check passed. Retrying immediately..." + recovery_stats.location_up = True Notifications.emit(EventType.STATUS_UPDATE, notification) logger.info(notification) continue # Proceed to attempt 2 @@ -360,11 +490,11 @@ def establish_system_connection( # ================================= # MTU # ================================= - should_test_mtu, current_mtu = should_we_even_test_mtu(profile.id) + should_test_mtu, current_mtu = should_we_even_test_mtu(profile_id=profile.id, recovery_stats=recovery_stats) if should_test_mtu: logger.info("We should test (and potentially change) the MTU..") - if was_mtu_changed(profile=profile, current_mtu=current_mtu): + if was_mtu_changed(profile=profile, current_mtu=current_mtu, recovery_stats=recovery_stats): logger.info("MTU was changed, so we're trying again.") continue else: @@ -382,11 +512,15 @@ def establish_system_connection( # ================================= # PORT # ================================= - if should_we_try_to_change_port(): + if should_we_try_to_change_port(recovery_stats): logger.info("We should try to change port") # Offer the choice to the end user: - return offer_port_change(profile) + message = get_port_change_message(profile=profile, recovery_stats=recovery_stats) + yes_choice = partial(change_port_then_reconnect, profile=profile, ignore=ignore, connection_observer=connection_observer, ticket_observer=ticket_observer) + # recovery_stats.stashed_partial = yes_choice + return Result(valid=False, function_choice=yes_choice, error_type=ResultError.PORT, message=message) + else: logger.info("Skipping port choice as an offer.") @@ -394,41 +528,22 @@ def establish_system_connection( # WIREGUARD RENEGOTIATION # ================================= terminate_system_connection(firewall_setting, dns_setting) - func_choice = partial(register_wireguard_session, profile=profile, connection_observer=connection_observer) - return Result(valid=False, function_choice=func_choice, error_type=ResultError.RENEGOTIATION, message="Would you like to try Wireguard key renegotiation? (Note: this invalidates other profiles with the same subscription.)") + + if not recovery_stats.asked_to_change_wg_key: + recovery_stats.asked_to_change_wg_key = True + yes_option = partial(renegotiate_and_reconnect, profile=profile, ignore=ignore, connection_observer=connection_observer, ticket_observer=ticket_observer) + no_option = partial(establish_system_connection, profile=profile, ignore=ignore, connection_observer=connection_observer, ticket_observer=ticket_observer) + return Result(valid=False, function_choice=yes_option, no_option=no_option, error_type=ResultError.RENEGOTIATION, message="Would you like to try Wireguard key renegotiation? (Note: this invalidates other profiles with the same subscription.)") + else: + logger.info("entering final loop") + return _conclude_loop_failure() - - - - - - - -# except (ConnectionError, CalledProcessError) as e: - -# # single time port is offered forever. -# if port_was_changed_already_today(): -# # permanent config modify: -# ConfigurationController.change_port_prompts(False) - -# else: -# notification = "Failed to connect." -# Notifications.emit(EventType.STATUS_UPDATE, notification) -# _cleanup_on_error(firewall_setting, dns_setting) -# raise ConnectionError("Failed to connect.") - -# else: -# # Attempt 3 failed -# logger.error("Maximum retries reached. Connection failed.") -# notification = "Failed to connect." -# Notifications.emit(EventType.STATUS_UPDATE, notification) -# _cleanup_on_error(firewall_setting, dns_setting) -# raise ConnectionError(notification) -# # return Result(valid=False, error_type=ResultError.CONNECTION, message="Failed to connect.") - -# First take -# _cleanup_on_error(firewall_setting, dns_setting) -# _cleanup_on_error(firewall_setting, dns_setting) -# # return Result(valid=False, error_type=ResultError.CONNECTION, message="Failed to connect.") +# ============= UPDATE STATE ============= +# if dns_setting and not did_dns_work: +# current_state.firewalled = firewall_tracker +# current_state.dns_set = did_dns_work +# SystemStateController.update_or_create(current_state) +# notification = "DNS Set Did Not Work." +# Notifications.emit(EventType.STATUS_UPDATE, notification) diff --git a/core/services/payment_phase/payment_state.py b/core/services/payment_phase/payment_state.py index 97a866a..2994ce9 100644 --- a/core/services/payment_phase/payment_state.py +++ b/core/services/payment_phase/payment_state.py @@ -1,147 +1,148 @@ -from core.models.manage.validate_dataclass import ValidatedDataclass -from core.utils.basic_operations.file_tools import delete_file +from core.utils.singleton_state import singleton, StatePersistenceMixin +# from core.models.manage.validate_dataclass import ValidatedDataclass +# from core.utils.basic_operations.file_tools import delete_file from core.Constants import Constants from core.models.enum_models import Currency, ProtocolChoice, PaymentType, PaymentCategory from core.errors.logger import logger -from dataclasses import dataclass, fields -import json +from dataclasses import dataclass #, fields +# import json import os -from functools import wraps -from enum import Enum +# from functools import wraps +# from enum import Enum from typing import Any, Optional -def singleton(cls): - instances = {} +# def singleton(cls): +# instances = {} - @wraps(cls) - def get_instance(*args, **kwargs): - instance_id = kwargs.get('instance_id') +# @wraps(cls) +# def get_instance(*args, **kwargs): +# instance_id = kwargs.get('instance_id') - # If instance_id provided, use localized mode (one instance per ID) - # Otherwise, use global mode (single instance) - key = (cls, instance_id) if instance_id else cls +# # If instance_id provided, use localized mode (one instance per ID) +# # Otherwise, use global mode (single instance) +# key = (cls, instance_id) if instance_id else cls - if key not in instances: - print(f"Creating new instance with key: {key}") - instances[key] = cls(*args, **kwargs) - else: - print(f"Reusing existing instance with key: {key}") +# if key not in instances: +# print(f"Creating new instance with key: {key}") +# instances[key] = cls(*args, **kwargs) +# else: +# print(f"Reusing existing instance with key: {key}") - # either way: - return instances[key] +# # either way: +# return instances[key] - return get_instance +# return get_instance -# for saving. -EXCLUDED_FIELDS = {"parent_folder", "folder_path", "_data_file"} +# # for saving. +# EXCLUDED_FIELDS = {"parent_folder", "folder_path", "_data_file"} -class StatePersistenceMixin(ValidatedDataclass): - def _serialize_value(self, value): - """Helper method to ensure values are JSON-serializable""" - if isinstance(value, Enum): - return value.value - return value +# class StatePersistenceMixin(ValidatedDataclass): +# def _serialize_value(self, value): +# """Helper method to ensure values are JSON-serializable""" +# if isinstance(value, Enum): +# return value.value +# return value - def prep_payload(self) -> dict|None: - """Orchestrate Save, Validation, and Payload preparation of relevant fields""" +# def prep_payload(self) -> dict|None: +# """Orchestrate Save, Validation, and Payload preparation of relevant fields""" - if not self.save(): - logger.error("Failed to Save! Not returning a payload") - return None +# if not self.save(): +# logger.error("Failed to Save! Not returning a payload") +# return None - payload = {} +# payload = {} - for f in fields(self): - # Only add the field if it's in our allowed list - if f.name in self.PAYLOAD_FIELDS: - # Get the raw value - raw_value = getattr(self, f.name) - # Clean it (handle Enums) - clean_value = self._serialize_value(raw_value) - # Add it to the dictionary - payload[f.name] = clean_value +# for f in fields(self): +# # Only add the field if it's in our allowed list +# if f.name in self.PAYLOAD_FIELDS: +# # Get the raw value +# raw_value = getattr(self, f.name) +# # Clean it (handle Enums) +# clean_value = self._serialize_value(raw_value) +# # Add it to the dictionary +# payload[f.name] = clean_value - return payload +# return payload - def save(self) -> bool: - """Save state to file""" +# def save(self) -> bool: +# """Save state to file""" - if not self.validate_payload(): - return False +# if not self.validate_payload(): +# return False - os.makedirs(self.folder_path, exist_ok=True) +# os.makedirs(self.folder_path, exist_ok=True) - # 1. Start with an empty dictionary - data = {} +# # 1. Start with an empty dictionary +# data = {} - # 2. Loop through the fields - for f in fields(self): - # Only add the field if it's NOT in the excluded list - if f.name not in EXCLUDED_FIELDS: - # Get the raw value - raw_value = getattr(self, f.name) - # Clean it (handle Enums) - clean_value = self._serialize_value(raw_value) - # Add it to the dictionary - data[f.name] = clean_value +# # 2. Loop through the fields +# for f in fields(self): +# # Only add the field if it's NOT in the excluded list +# if f.name not in EXCLUDED_FIELDS: +# # Get the raw value +# raw_value = getattr(self, f.name) +# # Clean it (handle Enums) +# clean_value = self._serialize_value(raw_value) +# # Add it to the dictionary +# data[f.name] = clean_value - print(f"saving to {self._data_file}") - with open(self._data_file, "w") as f: - json.dump(data, f, indent=4, sort_keys=True) +# print(f"saving to {self._data_file}") +# with open(self._data_file, "w") as f: +# json.dump(data, f, indent=4, sort_keys=True) - return True +# return True - def load(self): - """Load state from file, if each item passes a type check. With Enum translation""" +# def load(self): +# """Load state from file, if each item passes a type check. With Enum translation""" - logger.info("Starting load of pre-existing data") +# logger.info("Starting load of pre-existing data") - if not os.path.exists(self._data_file): - logger.info("does NOT exist already") - return +# if not os.path.exists(self._data_file): +# logger.info("does NOT exist already") +# return - with open(self._data_file, "r") as f: - data = json.load(f) +# with open(self._data_file, "r") as f: +# data = json.load(f) - for field in fields(self): - value = data.get(field.name) - logger.info(f"testing {field.name} of {value}") - if value is None: - continue +# for field in fields(self): +# value = data.get(field.name) +# logger.info(f"testing {field.name} of {value}") +# if value is None: +# continue - # Translation: If the expected type is an Enum, convert the raw value back into an Enum member - if isinstance(field.type, type) and issubclass(field.type, Enum): - logger.info("testing enum") - try: - value = field.type(value) - except ValueError: - logger.error(f"Invalid enum value {value} for {field.name}") - continue +# # Translation: If the expected type is an Enum, convert the raw value back into an Enum member +# if isinstance(field.type, type) and issubclass(field.type, Enum): +# logger.info("testing enum") +# try: +# value = field.type(value) +# except ValueError: +# logger.error(f"Invalid enum value {value} for {field.name}") +# continue - # type check works because 'value' is now an actual Enum member - if isinstance(value, field.type): - setattr(self, field.name, value) +# # type check works because 'value' is now an actual Enum member +# if isinstance(value, field.type): +# setattr(self, field.name, value) - def wipe_data_on_disk(self): - logger.info("Starting wipe of pre-existing data") +# def wipe_data_on_disk(self): +# logger.info("Starting wipe of pre-existing data") - if not os.path.exists(self._data_file): - logger.info("File does NOT exist already") - return +# if not os.path.exists(self._data_file): +# logger.info("File does NOT exist already") +# return - return delete_file(self._data_file) +# return delete_file(self._data_file) - def exists_already(self) -> bool: - return os.path.exists(self._data_file) +# def exists_already(self) -> bool: +# return os.path.exists(self._data_file) HOW_MANY_PROFILES_DEFAULT = 6 @@ -177,13 +178,14 @@ class TicketPaymentData(StatePersistenceMixin): def wipe(self): """Wipe all data and start over.""" wiped = self.wipe_data_on_disk() + logger.info(f"Wiped: {wiped}") self.currency = None self.which_key = None self.how_many_profiles = 6 self.billing_code = None self.use_random = False self.paid = False - sefl.billing_pending = True + self.billing_pending = True @singleton @@ -218,6 +220,7 @@ class PaymentData(StatePersistenceMixin): def wipe(self): """Wipe all data and start over.""" wiped = self.wipe_data_on_disk() + logger.info(f"Wiped: {wiped}") self.currency = None self.payment_type = None self.plan_id = None diff --git a/core/utils/singleton.py b/core/utils/singleton.py deleted file mode 100644 index ca6f038..0000000 --- a/core/utils/singleton.py +++ /dev/null @@ -1,25 +0,0 @@ - -from functools import wraps - -def singleton(cls): - instances = {} - - @wraps(cls) - def get_instance(*args, **kwargs): - instance_id = kwargs.get('instance_id') - - # If instance_id provided, use localized mode (one instance per ID) - # Otherwise, use global mode (single instance) - key = (cls, instance_id) if instance_id else cls - - if key not in instances: - print(f"Creating new instance with key: {key}") - instances[key] = cls(*args, **kwargs) - else: - print(f"Reusing existing instance with key: {key}") - - # either way: - return instances[key] - - return get_instance - diff --git a/core/utils/singleton_state.py b/core/utils/singleton_state.py new file mode 100644 index 0000000..27d2de2 --- /dev/null +++ b/core/utils/singleton_state.py @@ -0,0 +1,161 @@ +from core.models.manage.validate_dataclass import ValidatedDataclass +from core.utils.basic_operations.file_tools import delete_file +from core.errors.logger import logger + +from dataclasses import fields, MISSING +import json +import os +from functools import wraps +from enum import Enum +from typing import Any + + +def singleton(cls): + instances = {} + + @wraps(cls) + def get_instance(*args, **kwargs): + instance_id = kwargs.get('instance_id') + + # If instance_id provided, use localized mode (one instance per ID) + # Otherwise, use global mode (single instance) + key = (cls, instance_id) if instance_id else cls + + if key not in instances: + print(f"Creating new instance with key: {key}") + instances[key] = cls(*args, **kwargs) + else: + print(f"Reusing existing instance with key: {key}") + + # either way: + return instances[key] + + return get_instance + + +# for saving. +EXCLUDED_FIELDS = {"parent_folder", "folder_path", "_data_file"} + + +class StatePersistenceMixin(ValidatedDataclass): + def _serialize_value(self, value): + """Helper method to ensure values are JSON-serializable""" + if isinstance(value, Enum): + return value.value + return value + + + def prep_payload(self) -> dict|None: + """Orchestrate Save, Validation, and Payload preparation of relevant fields""" + + if not self.save(): + logger.error("Failed to Save! Not returning a payload") + return None + + payload = {} + + for f in fields(self): + # Only add the field if it's in our allowed list + if f.name in self.PAYLOAD_FIELDS: + # Get the raw value + raw_value = getattr(self, f.name) + # Clean it (handle Enums) + clean_value = self._serialize_value(raw_value) + # Add it to the dictionary + payload[f.name] = clean_value + + return payload + + + def save(self) -> bool: + """Save state to file""" + + if not self.validate_payload(): + return False + + os.makedirs(self.folder_path, exist_ok=True) + + # 1. Start with an empty dictionary + data = {} + + # 2. Loop through the fields + for f in fields(self): + # Only add the field if it's NOT in the excluded list + if f.name not in EXCLUDED_FIELDS: + # Get the raw value + raw_value = getattr(self, f.name) + # Clean it (handle Enums) + clean_value = self._serialize_value(raw_value) + # Add it to the dictionary + data[f.name] = clean_value + + print(f"saving to {self._data_file}") + with open(self._data_file, "w") as f: + json.dump(data, f, indent=4, sort_keys=True) + + return True + + + def load(self): + """Load state from file, if each item passes a type check. With Enum translation""" + + logger.info("Starting load of pre-existing data") + + if not os.path.exists(self._data_file): + logger.info("does NOT exist already") + return + + with open(self._data_file, "r") as f: + data = json.load(f) + + for field in fields(self): + value = data.get(field.name) + logger.info(f"testing {field.name} of {value}") + if value is None: + continue + + # Translation: If the expected type is an Enum, convert the raw value back into an Enum member + if isinstance(field.type, type) and issubclass(field.type, Enum): + logger.info("testing enum") + try: + value = field.type(value) + except ValueError: + logger.error(f"Invalid enum value {value} for {field.name}") + continue + + # type check works because 'value' is now an actual Enum member + if isinstance(value, field.type): + setattr(self, field.name, value) + + + def wipe_data_on_disk(self): + logger.info("Starting wipe of pre-existing data") + + if not os.path.exists(self._data_file): + logger.info("File does NOT exist already") + return + + return delete_file(self._data_file) + + + def exists_already(self) -> bool: + return os.path.exists(self._data_file) + + + +def wipe_model(model: Any, KEEP_FIELDS: list) -> Any: + for field in fields(model): + if field.name not in KEEP_FIELDS: + # 1. Check for a static default value (e.g., attempts: int = 0) + if field.default is not MISSING: + setattr(model, field.name, field.default) + + # 2. Check for a default factory (e.g., items: list = field(default_factory=list)) + elif field.default_factory is not MISSING: + setattr(model, field.name, field.default_factory()) + + # 3. Fallback: if no default is defined at all, set it to None. + else: + setattr(model, field.name, None) + + return model