sp-hydra-veil-core/core/controllers/ConnectionController.py

from core.Constants import Constants
from core.Errors import InvalidSubscriptionError, MissingSubscriptionError, ConnectionUnprotectedError, ConnectionTerminationError
from core.controllers.ConfigurationController import ConfigurationController
from core.controllers.ProfileController import ProfileController
from core.controllers.SessionStateController import SessionStateController
from core.controllers.SystemStateController import SystemStateController
from core.models.session.SessionProfile import SessionProfile
from core.models.system.SystemProfile import SystemProfile
from core.models.system.SystemState import SystemState
from core.observers import ConnectionObserver
from core.services.WebServiceApiService import WebServiceApiService
from pathlib import Path
from typing import Union, Optional
import os
import re
import requests
import socket
import subprocess
import tempfile
import time


class ConnectionController:

    @staticmethod
    def with_preferred_connection(*args, task: callable, connection_observer: Optional[ConnectionObserver] = None, **kwargs):

        connection = ConfigurationController.get_connection()

        if connection == 'system':
            return task(*args, **kwargs)

        elif connection == 'tor':
            return ConnectionController.__with_tor_connection(*args, task=task, connection_observer=connection_observer, **kwargs)

    @staticmethod
    def establish_connection(profile: Union[SessionProfile, SystemProfile], force: bool = False, connection_observer: Optional[ConnectionObserver] = None):

        connection = profile.connection

        if connection.needs_proxy_configuration() and not profile.has_proxy_configuration():

            if profile.has_subscription():

                if not profile.subscription.has_been_activated():
                    ProfileController.activate_subscription(profile, connection_observer=connection_observer)

                proxy_configuration = ConnectionController.with_preferred_connection(profile.subscription.billing_code, task=WebServiceApiService.get_proxy_configuration, connection_observer=connection_observer)

                if proxy_configuration is None:
                    raise InvalidSubscriptionError()

                profile.attach_proxy_configuration(proxy_configuration)

            else:
                raise MissingSubscriptionError()

        if connection.needs_wireguard_configuration() and not profile.has_wireguard_configuration():

            if profile.has_subscription():

                if not profile.subscription.has_been_activated():
                    ProfileController.activate_subscription(profile, connection_observer=connection_observer)

                ProfileController.register_wireguard_session(profile, connection_observer=connection_observer)

            else:

                if profile.is_system_profile():

                    if ConnectionController.system_uses_wireguard_interface() and SystemStateController.exists():

                        active_profile = ProfileController.get(SystemStateController.get().profile_id)

                        try:
                            ConnectionController.terminate_system_connection(active_profile)
                        except ConnectionTerminationError:
                            pass

                raise MissingSubscriptionError()

        if profile.is_session_profile():

            try:
                return ConnectionController.establish_session_connection(profile, force=force, connection_observer=connection_observer)

            except ConnectionError:

                if ConnectionController.__should_renegotiate(profile):

                    ProfileController.register_wireguard_session(profile, connection_observer=connection_observer)
                    return ConnectionController.establish_session_connection(profile, force=force, connection_observer=connection_observer)

                else:
                    raise ConnectionError('The connection could not be established.')

        if profile.is_system_profile():

            try:
                return ConnectionController.establish_system_connection(profile, connection_observer=connection_observer)

            except ConnectionError:

                if ConnectionController.__should_renegotiate(profile):

                    ProfileController.register_wireguard_session(profile, connection_observer=connection_observer)
                    return ConnectionController.establish_system_connection(profile, connection_observer=connection_observer)

                else:
                    raise ConnectionError('The connection could not be established.')

    @staticmethod
    def establish_session_connection(profile: SessionProfile, force: Optional[bool] = False, connection_observer: Optional[ConnectionObserver] = None):

        session_directory = tempfile.mkdtemp(prefix='sp-')
        session_state = SessionStateController.get_or_new(profile.id)

        port_number = None
        proxy_port_number = None

        if profile.connection.is_unprotected():

            if not ConnectionController.system_uses_wireguard_interface():

                if not force:
                    raise ConnectionUnprotectedError('Connection unprotected while the system is not using a WireGuard interface.')
                else:
                    ProfileController.disable(profile)

        if profile.connection.code == 'tor':

            port_number = ConnectionController.get_random_available_port_number()
            ConnectionController.establish_tor_session_connection(session_directory, port_number)
            session_state.network_port_numbers.append(port_number)

        elif profile.connection.code == 'wireguard':

            port_number = ConnectionController.get_random_available_port_number()
            ConnectionController.establish_wireguard_session_connection(profile, session_directory, port_number)
            session_state.network_port_numbers.append(port_number)

        if profile.connection.masked:

            while proxy_port_number is None or proxy_port_number == port_number:
                proxy_port_number = ConnectionController.get_random_available_port_number()

            ConnectionController.establish_proxy_session_connection(profile, session_directory, port_number, proxy_port_number)
            session_state.network_port_numbers.append(proxy_port_number)

        if not profile.connection.is_unprotected():
            ConnectionController.await_connection(proxy_port_number or port_number, connection_observer=connection_observer)

        SessionStateController.update_or_create(session_state)

        return proxy_port_number or port_number

    @staticmethod
    def establish_system_connection(profile: SystemProfile, connection_observer: Optional[ConnectionObserver] = None):

        if subprocess.getstatusoutput('pkexec --help')[0] == 127:
            raise OSError('The polkit toolkit does not appear to be installed.')

        if subprocess.getstatusoutput('wg-quick --help')[0] == 127:
            raise OSError('WireGuard tools does not appear to be installed.')

        process = subprocess.Popen(('pkexec', 'wg-quick', 'up', profile.get_wireguard_configuration_path()), stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
        completed_successfully = not bool(os.waitpid(process.pid, 0)[1] >> 8)

        if completed_successfully:

            SystemStateController.update_or_create(SystemState(profile.id))

            try:
                ConnectionController.await_connection(connection_observer=connection_observer)

            except ConnectionError:

                ConnectionController.terminate_system_connection(profile)
                raise ConnectionError('The connection could not be established.')

        else:

            ConnectionController.terminate_system_connection(profile)

            process = subprocess.Popen(('pkexec', 'wg-quick', 'up', profile.get_wireguard_configuration_path()), stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
            completed_successfully = not bool(os.waitpid(process.pid, 0)[1] >> 8)

            if completed_successfully:

                SystemStateController.update_or_create(SystemState(profile.id))

                try:
                    ConnectionController.await_connection(connection_observer=connection_observer)

                except ConnectionError:

                    ConnectionController.terminate_system_connection(profile)
                    raise ConnectionError('The connection could not be established.')

            else:
                raise ConnectionError('The connection could not be established.')

        time.sleep(1.0)

    @staticmethod
    def establish_tor_session_connection(session_directory: str, port_number: int):

        if subprocess.getstatusoutput('tor --help')[0] == 127:
            raise OSError('Tor does not appear to be installed.')

        tor_session_directory = f'{session_directory}/tor'
        Path(tor_session_directory).mkdir(exist_ok=True, mode=0o700)

        process = subprocess.Popen(('echo', f'DataDirectory {tor_session_directory}/tor\nSocksPort {port_number}'), stdout=subprocess.PIPE)
        return subprocess.Popen(('tor', '-f', '-'), stdin=process.stdout, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)

    @staticmethod
    def establish_wireguard_session_connection(profile: SessionProfile, session_directory: str, port_number: int):

        if not profile.has_wireguard_configuration():
            raise FileNotFoundError('No valid WireGuard configuration file detected.')

        wireguard_session_directory = f'{session_directory}/wireguard'
        Path(wireguard_session_directory).mkdir(exist_ok=True, mode=0o700)

        wireproxy_configuration_file_path = f'{wireguard_session_directory}/wireproxy.conf'
        Path(wireproxy_configuration_file_path).touch(exist_ok=True, mode=0o600)

        with open(wireproxy_configuration_file_path, 'w') as wireproxy_configuration_file:

            wireproxy_configuration_file.write(f'WGConfig = {profile.get_wireguard_configuration_path()}\n\n[Socks5]\nBindAddress = 127.0.0.1:{str(port_number)}\n')
            wireproxy_configuration_file.close()

        return subprocess.Popen((f'{Constants.SP_DATA_HOME}/wireproxy/wireproxy', '-c', wireproxy_configuration_file_path), stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)

    @staticmethod
    def establish_proxy_session_connection(profile: SessionProfile, session_directory: str, port_number: int, proxy_port_number: int):

        if subprocess.getstatusoutput('proxychains4 --help')[0] == 127:
            raise OSError('ProxyChains does not appear to be installed.')

        if subprocess.getstatusoutput('microsocks --help')[0] == 127:
            raise OSError('MicroSocks does not appear to be installed.')

        if profile.has_proxy_configuration():
            proxy_configuration = profile.get_proxy_configuration()
        else:
            raise FileNotFoundError('No valid proxy configuration file detected.')

        proxy_session_directory = f'{session_directory}/proxy'
        Path(proxy_session_directory).mkdir(parents=True, exist_ok=True, mode=0o700)

        proxychains_proxy_list = ''

        if port_number is not None:
            proxychains_proxy_list = f'socks5 127.0.0.1 {port_number}\n'

        proxychains_proxy_list = proxychains_proxy_list + f'socks5 {proxy_configuration.ip_address} {proxy_configuration.port_number} {proxy_configuration.username} {proxy_configuration.password}'
        proxychains_template_file_path = f'{Constants.SP_DATA_HOME}/proxychains.ptpl'

        with open(proxychains_template_file_path, 'r') as proxychains_template_file:

            proxychains_configuration_file_path = f'{proxy_session_directory}/proxychains.conf'
            Path(proxychains_configuration_file_path).touch(exist_ok=True, mode=0o600)

            proxychains_configuration_file_contents = proxychains_template_file.read().format(proxy_list=proxychains_proxy_list)

            with open(proxychains_configuration_file_path, 'w') as proxychains_configuration_file:

                proxychains_configuration_file.write(proxychains_configuration_file_contents)
                proxychains_configuration_file.close()

        return subprocess.Popen(('proxychains4', '-f', proxychains_configuration_file_path, 'microsocks', '-p', str(proxy_port_number)), stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)

    @staticmethod
    def terminate_system_connection(profile: SystemProfile):

        if subprocess.getstatusoutput('pkexec --help')[0] == 127:
            raise OSError('The polkit toolkit does not appear to be installed.')

        if subprocess.getstatusoutput('wg-quick --help')[0] == 127:
            raise OSError('WireGuard tools does not appear to be installed.')

        process = subprocess.Popen(('pkexec', 'wg-quick', 'down', profile.get_wireguard_configuration_path()), stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
        completed_successfully = not bool(os.waitpid(process.pid, 0)[1] >> 8)

        if completed_successfully or not ConnectionController.system_uses_wireguard_interface():

            system_state = SystemStateController.get()

            if system_state is not None:
                system_state.dissolve()

        else:
            raise ConnectionTerminationError('The connection could not be terminated.')

    @staticmethod
    def get_proxies(port_number: int):

        return dict(
            http=f'socks5h://127.0.0.1:{port_number}',
            https=f'socks5h://127.0.0.1:{port_number}'
        )

    @staticmethod
    def get_random_available_port_number():

        socket_instance = socket.socket()
        socket_instance.bind(('', 0))
        port_number = socket_instance.getsockname()[1]
        socket_instance.close()

        return port_number

    @staticmethod
    def await_connection(port_number: int = None, connection_observer: Optional[ConnectionObserver] = None):

        if port_number is None:
            ConnectionController.await_network_interface()

        maximum_number_of_attempts = 5
        retry_interval = 5.0

        for retry_count in range(maximum_number_of_attempts):

            if connection_observer is not None:

                connection_observer.notify('connecting', dict(
                    retry_interval=retry_interval,
                    maximum_number_of_attempts=maximum_number_of_attempts,
                    attempt_count=retry_count + 1
                ))

            try:

                ConnectionController.__test_connection(port_number)
                return

            except ConnectionError:

                time.sleep(retry_interval)
                retry_count += 1

        raise ConnectionError('The connection could not be established.')

    @staticmethod
    def await_network_interface():

        network_interface_is_activated = False

        retry_interval = .5
        maximum_number_of_attempts = 10
        attempt = 0

        while not network_interface_is_activated and attempt < maximum_number_of_attempts:

            time.sleep(retry_interval)

            network_interface_is_activated = ConnectionController.system_uses_wireguard_interface()
            attempt += 1

        if network_interface_is_activated is False:
            ConnectionError('The network interface could not be activated.')

    @staticmethod
    def system_uses_wireguard_interface():

        if subprocess.getstatusoutput('ip --help')[0] == 127:
            raise OSError('The iproute2 utility package does not appear to be installed.')

        process = subprocess.Popen(('ip', 'route', 'get', '192.0.2.1'), stdout=subprocess.PIPE)
        process_output = str(process.stdout.read())

        return bool(re.search('dev wg', str(process_output)))

    @staticmethod
    def __with_tor_connection(*args, task: callable, connection_observer: Optional[ConnectionObserver] = None, **kwargs):

        session_directory = tempfile.mkdtemp(prefix='sp-')
        port_number = ConnectionController.get_random_available_port_number()
        process = ConnectionController.establish_tor_session_connection(session_directory, port_number)

        ConnectionController.await_connection(port_number, connection_observer=connection_observer)
        task_output = task(*args, proxies=ConnectionController.get_proxies(port_number), **kwargs)
        process.terminate()

        return task_output

    @staticmethod
    def __test_connection(port_number: Optional[int] = None, timeout: float = 10.0):

        timeout = float(timeout)
        proxies = None

        if port_number is not None:
            proxies = ConnectionController.get_proxies(port_number)

        try:
            requests.get(f'{Constants.SP_API_BASE_URL}/health', timeout=timeout, proxies=proxies)
        except requests.exceptions.RequestException:
            raise ConnectionError('The connection could not be established.')

    @staticmethod
    def __should_renegotiate(profile: Union[SessionProfile, SystemProfile]):

        if profile.connection.needs_wireguard_configuration() and profile.has_wireguard_configuration():

            if profile.subscription.has_been_activated():
                return True

        return False