hydraveil-gui/gui/v2/workers/worker.py

from datetime import datetime, timezone

from PyQt6.QtCore import QObject, pyqtSignal

from core.controllers.ProfileController import ProfileController
from core.controllers.SubscriptionController import SubscriptionController
from core.controllers.tickets.UseTicketController import (
    use_ticket,
    do_we_use_a_random_ticket,
)
from core.models.session.SessionProfile import SessionProfile
from core.models.system.SystemProfile import SystemProfile
from core.models.Result import Result, ResultError
from core.errors.exceptions import SudoScript, MissingPreReqs, FirewallError
from core.Errors import (
    CommandNotFoundError,
    EndpointVerificationError,
    FileIntegrityError,
    InvalidSubscriptionError,
    MissingSubscriptionError,
    ProfileActivationError,
    ProfileModificationError,
    ProfileStateConflictError,
    UnsupportedApplicationVersionError,
)

from gui.v2.actions.database_health import GuiStorageDatabaseError
from gui.v2.actions.operation_result_dispatch import dispatch_result_action
from gui.v2.actions.operation_results import result_from_exception
from gui.v2.infrastructure.screen_size import get_max_screensize
from gui.v2.infrastructure.setup_observers import (
    application_version_observer,
    connection_observer,
    profile_observer,
    ticket_observer,
)


class Worker(QObject):
    update_signal = pyqtSignal(str, bool, int, int, str)
    change_page = pyqtSignal(str, bool)
    ticket_data_loss = pyqtSignal(str, str)
    operation_failed = pyqtSignal(object)

    def __init__(self, profile_data):
        self.profile_data = profile_data
        super().__init__()
        profile_observer.subscribe(
            'disabled', lambda event: self.handle_profile_status(event.subject, False))
        profile_observer.subscribe(
            'enabled', lambda event: self.handle_profile_status(event.subject, True))
        self.profile_type = None
        self._ticket_error_emitted = False
        self._consumed_ticket = None
        self._pending_operation_result = None
        self._pending_operation_kwargs = {}

    def run(self):
        try:
            self.profile = ProfileController.get(int(self.profile_data['id']))
        except GuiStorageDatabaseError:
            self.update_signal.emit(
                "Local storage database could not be read. Restart and recover storage.db.", False, None, None, None)
            return
        except Exception:
            self.update_signal.emit(
                "Could not load profile data. Sync or restart and try again.", False, None, None, None)
            return

        incomplete_reason = self._profile_incomplete_reason()
        if incomplete_reason:
            self.update_signal.emit(incomplete_reason, False, None, None, None)
            return

        if 'use_ticket' in self.profile_data:
            ticket_billing_code = self._consume_ticket(
                self.profile_data['use_ticket'])
            if ticket_billing_code is None:
                return
            self.profile_data['billing_code'] = ticket_billing_code
        elif 'billing_code' not in self.profile_data and self.profile is not None:
            if not self._profile_has_valid_subscription():
                ticket_billing_code = self._maybe_auto_use_ticket()
                if ticket_billing_code is None and self._ticket_error_emitted:
                    return
                if ticket_billing_code:
                    self.profile_data['billing_code'] = ticket_billing_code

        if 'billing_code' in self.profile_data:
            try:
                subscription = SubscriptionController.get(
                    self.profile_data['billing_code'], connection_observer=connection_observer)
            except Exception as e:
                if self._consumed_ticket is not None:
                    self._emit_subscription_lookup_failure(exception=e)
                    return
                subscription = None
            if subscription is not None:
                ProfileController.attach_subscription(
                    self.profile, subscription)
            else:
                if self._consumed_ticket is not None:
                    self._emit_subscription_lookup_failure()
                    return
                self.change_page.emit('The billing code is invalid.', True)
                return
        if self.profile:
            try:
                ignore_exceptions = []
                if self.profile_data.get('ignore_endpoint_verification', False):
                    ignore_exceptions.append(EndpointVerificationError)
                if self.profile_data.get('ignore_profile_state_conflict', False):
                    ignore_exceptions.append(ProfileStateConflictError)
                ignore_tuple = tuple(ignore_exceptions)
                max_resolution = get_max_screensize()
                enable_result = ProfileController.enable(self.profile, ignore=ignore_tuple, profile_observer=profile_observer,
                                                         application_version_observer=application_version_observer,
                                                         connection_observer=connection_observer, ticket_observer=ticket_observer, max_resolution=max_resolution)
                if isinstance(enable_result, Result) and not enable_result.valid:
                    self._emit_operation_failure(enable_result)
                    return
            except EndpointVerificationError:
                self.update_signal.emit(
                    "ENDPOINT_VERIFICATION_ERROR", False, self.profile_data['id'], None, None)
            except (InvalidSubscriptionError, MissingSubscriptionError) as e:
                if self._consumed_ticket is not None:
                    self._emit_subscription_lookup_failure(exception=e)
                else:
                    self.change_page.emit(
                        f"Subscription missing or invalid for profile {self.profile_data['id']}", True)
            except ProfileActivationError:
                self.update_signal.emit(
                    "The profile could not be enabled", False, None, None, None)
            except UnsupportedApplicationVersionError:
                self.update_signal.emit(
                    "The application version in question is not supported", False, None, None, None)
            except FileIntegrityError:
                self.update_signal.emit(
                    "Application version file integrity could not be verified.", False, None, None, None)
            except ProfileModificationError:
                self.update_signal.emit(
                    "WireGuard configuration could not be attached.", False, None, None, None)
            except ProfileStateConflictError:
                self.update_signal.emit(
                    "PROFILE_STATE_CONFLICT_ERROR", False, self.profile_data['id'], None, None)
            except CommandNotFoundError as e:
                self.update_signal.emit(str(e.subject), False, -1, None, None)
            except SudoScript as e:
                self.update_signal.emit(str(e), False, None, None, None)
            except MissingPreReqs as e:
                self.update_signal.emit(str(e), False, None, None, None)
            except FirewallError as e:
                self.update_signal.emit(str(e), False, None, None, None)
            except Exception as e:
                print(e)
                self.update_signal.emit(
                    "An unknown error occurred", False, None, None, None)

        else:
            self.update_signal.emit(
                f"No profile found with ID: {self.profile_data['id']}", False, None, None, None)

    def _profile_incomplete_reason(self):
        if self.profile is None:
            return None

        connection = getattr(self.profile, 'connection', None)
        if not getattr(connection, 'code', None):
            return "Profile connection data is incomplete. Sync the database and try again."

        location = getattr(self.profile, 'location', None)
        if location is None or isinstance(location, dict):
            return "Profile location data is incomplete. Sync the database and try again."

        if isinstance(self.profile, SessionProfile):
            application_version = getattr(self.profile, 'application_version', None)
            if isinstance(application_version, dict) or application_version is None:
                return "Profile browser data is incomplete. Sync the database and try again."
            if not getattr(application_version, 'application_code', None) or not getattr(application_version, 'version_number', None):
                return "Profile browser data is incomplete. Sync the database and try again."

        return None

    def _profile_has_valid_subscription(self):
        try:
            sub = getattr(self.profile, 'subscription', None)
            if not sub:
                return False
            expires_at = getattr(sub, 'expires_at', None)
            if not expires_at:
                return False
            if expires_at.tzinfo is None:
                expires_at = expires_at.replace(tzinfo=timezone.utc)
            return expires_at > datetime.now(timezone.utc)
        except Exception:
            return False

    def _maybe_auto_use_ticket(self):
        profile_connection_type = self.profile.connection.code

        if profile_connection_type != "wireguard":
            error_msg = "Tickets are not yet supported for Tor.  We are very sorry, but we're not able to develop all features for all products."
            print(error_msg)
            return None

        if self.profile.assassin:
            print("skipping random ticket use for the assassin, as it already has it baked in.")
            return None
        try:
            print("using a random ticket from GUI...")
            which_ticket, error_msg = do_we_use_a_random_ticket(ticket_observer)
        except Exception as e:
            self._emit_ticket_failure(None, exception=e)
            return None
        if error_msg:
            self._emit_ticket_failure(
                None,
                result=Result(
                    valid=False,
                    error_type=ResultError.TICKET,
                    message=f'Ticket use failed: {error_msg}',
                ),
            )
            return None
        if which_ticket is None or which_ticket == 'error':
            return None
        return self._consume_ticket(which_ticket)

    def _ticket_location_id(self):
        location = getattr(self.profile, 'location', None)
        location_id = getattr(location, 'id', None)
        if location_id is None:
            return None
        if isinstance(location_id, str):
            location_id = location_id.strip()
        if location_id == '':
            return None
        try:
            return int(location_id)
        except (TypeError, ValueError):
            return None

    def _emit_ticket_failure(self, which_ticket, result=None, message=None, exception=None):
        self._ticket_error_emitted = True
        if exception is not None:
            result = result_from_exception(exception)
        elif result is None:
            result = Result(
                valid=False,
                error_type=ResultError.UNKNOWN,
                message=message or 'Ticket use failed.',
            )
        self._emit_operation_failure(result, which_ticket=which_ticket)

    def _emit_operation_failure(self, result, **kwargs):
        self._pending_operation_result = result
        self._pending_operation_kwargs = {
            key: value for key, value in kwargs.items()
            if value is not None
        }
        self.operation_failed.emit(result)

    def handle_operation_popup_choice(self, accepted):
        if not accepted:
            return None
        if not isinstance(self._pending_operation_result, Result):
            return None
        return dispatch_result_action(
            self._pending_operation_result,
            **self._pending_operation_kwargs,
        )

    def _emit_subscription_lookup_failure(self, exception=None):
        message = (
            "Ticket was accepted, but the subscription details could not be "
            "retrieved. Check your connection and try again."
        )
        if exception is not None:
            error_text = str(exception)
            if error_text:
                message = f"{message} {error_text}"
        self._emit_ticket_failure(
            self._consumed_ticket,
            result=Result(
                valid=False,
                error_type=ResultError.CONNECTION,
                message=message,
            ),
        )

    def _consume_ticket(self, which_ticket):
        which_location = self._ticket_location_id()
        if which_location is None:
            self._emit_ticket_failure(
                which_ticket,
                result=Result(
                    valid=False,
                    error_type=ResultError.MISSING_DATA,
                    message='Could not determine profile location for ticket use.',
                ),
            )
            return None

        try:
            outcome = use_ticket(
                which_ticket=which_ticket,
                which_location=which_location,
                ticket_observer=ticket_observer,
                connection_observer=connection_observer,
                profile=self.profile
            )
        except Exception as e:
            self._emit_ticket_failure(which_ticket, exception=e)
            return None

        if not isinstance(outcome, Result):
            self._emit_ticket_failure(
                which_ticket,
                result=Result(
                    valid=False,
                    error_type=ResultError.INVALID_API_REPLY,
                    message=f'use_ticket returned {type(outcome).__name__}, expected Result.',
                ),
            )
            return None

        if outcome.valid:
            billing_code = outcome.data
            if not billing_code:
                self._emit_ticket_failure(
                    which_ticket,
                    result=Result(
                        valid=False,
                        error_type=ResultError.INVALID_API_REPLY,
                        message='Ticket use succeeded but no billing code was returned.',
                    ),
                )
                return None
            self._consumed_ticket = str(which_ticket)
            return billing_code

        self._emit_ticket_failure(which_ticket, result=outcome)
        return None

    def handle_profile_status(self, profile, is_enabled):
        profile_id = profile.id
        profile_connection = str(profile.connection.code)
        message = self.generate_profile_message(profile, is_enabled)

        if isinstance(profile, SessionProfile):
            self.profile_type = 1
        elif isinstance(profile, SystemProfile):
            self.profile_type = 2
        else:
            self.profile_type = None

        self.update_signal.emit(
            message, is_enabled, profile_id, self.profile_type, profile_connection)

    @staticmethod
    def generate_profile_message(profile, is_enabled, idle=False):

        profile_id = profile.id
        if not profile.subscription or not profile.subscription.expires_at:
            return f"Offline. No subscription found."

        profile_date = profile.subscription.expires_at
        status = 'enabled' if is_enabled else 'disabled'

        expiration_date = profile_date.replace(tzinfo=timezone.utc)

        time_left = expiration_date - datetime.now(timezone.utc)
        days_left = time_left.days
        hours_left, remainder = divmod(time_left.seconds, 3600)

        formatted_expiration = expiration_date.strftime("%Y-%m-%d %H:%M:%S")

        if expiration_date < datetime.now(timezone.utc):
            return "Offline. Subscription has expired."
        if idle:
            return f"Offline. Expires in {days_left} days."

        if is_enabled:
            return f"Profile {int(profile_id)} {status}. Expires on {formatted_expiration}. Time left: {days_left} days, {hours_left} hours."
        else:
            return f"Profile {int(profile_id)} {status}"