sp-hydra-veil-core/core/models/system/SystemProfile.py

from core.Constants import Constants
from core.Errors import ProfileDeletionError, ProfileModificationError, CommandNotFoundError
from core.models.BaseProfile import BaseProfile
from core.models.system.SystemConnection import SystemConnection
from dataclasses import dataclass
from typing import Optional
import json
import os
import shutil
import subprocess

@dataclass
class SystemProfile(BaseProfile):
    connection: Optional[SystemConnection]

    def get_system_config_path(self):
        filepath = self.__get_system_config_path(self.id)
        the_id = self.id
        return filepath

    def save(self):
        if 'location' in self._get_dirty_keys():
            self.__delete_wireguard_configuration()
        super().save()

    def attach_wireguard_configuration(self, wireguard_configuration):
        wireguard_configuration_file_backup_path = f'{self.get_config_path()}/wg.conf.bak'
        with open(wireguard_configuration_file_backup_path, 'w') as wireguard_configuration_file:
            wireguard_configuration_file.write(wireguard_configuration)
        wireguard_configuration_is_attached = False
        failed_attempt_count = 0
        # Try sudo first (configured via sudoers by installer), fall back to pkexec
        install_cmd = 'sudo' if shutil.which('sudo') else 'pkexec'
        if install_cmd == 'pkexec' and shutil.which('pkexec') is None:
            raise CommandNotFoundError('pkexec')
        while not wireguard_configuration_is_attached and failed_attempt_count < 3:
            process = subprocess.Popen((install_cmd, 'install', '-D', wireguard_configuration_file_backup_path, self.get_wireguard_configuration_path(), '-o', 'root', '-m', '744'))
            wireguard_configuration_is_attached = not bool(os.waitpid(process.pid, 0)[1] >> 8)
            if not wireguard_configuration_is_attached:
                failed_attempt_count += 1
        if not wireguard_configuration_is_attached:
            raise ProfileModificationError('The WireGuard configuration could not be attached.')

    def get_wireguard_configuration_path(self):
        filepath = f'{self.get_system_config_path()}/wg.conf'
        return filepath

    def has_wireguard_configuration(self):
        filepath = f'{self.get_system_config_path()}/wg.conf'
        if os.path.isdir(os.path.dirname(filepath)):
            return os.path.isfile(filepath)
        else:
            return False

    def address_security_incident(self):
        super().address_security_incident()
        self.__delete_wireguard_configuration()

    def delete(self):
        try:
            self.__delete_wireguard_configuration()
        except ProfileModificationError:
            raise ProfileDeletionError('The WireGuard configuration could not be deleted.')
        if shutil.which('pkexec') is None:
            raise CommandNotFoundError('pkexec')

        try:
            process = subprocess.run(('pkexec', 'rm', '-rf', self.get_system_config_path()))
            completed_successfully = not bool(os.waitpid(process.pid, 0)[1] >> 8)
            if not completed_successfully:
                raise ProfileDeletionError('The profile could not be deleted.')
        except:
            print("skipping the delete of the WG folder.")

        super().delete()

    def attach_operator_proxy_session(self, operator_proxy_session):
        from core.models.OperatorProxySession import OperatorProxySession
        operator_proxy_session_file_contents = f'{operator_proxy_session.to_json(indent=4)}\n'
        os.makedirs(self.get_config_path(), exist_ok=True)
        operator_proxy_session_file_path = self.get_operator_proxy_session_path()
        with open(operator_proxy_session_file_path, 'w') as operator_proxy_session_file:
            operator_proxy_session_file.write(operator_proxy_session_file_contents)
        if operator_proxy_session.location_country_code and operator_proxy_session.location_city_code:
            try:
                from core.models.orm_models.Location import Location
                from core.models.manage.session_management import get_session
                from sqlalchemy import select
                session = get_session()
                loc = session.execute(
                    select(Location).where(
                        (Location.country_code == operator_proxy_session.location_country_code) &
                        (Location.code == operator_proxy_session.location_city_code)
                    )
                ).scalar_one_or_none()
                if loc:
                    self.location = loc
                    self.save()
            except Exception:
                pass

    def get_operator_proxy_session_path(self):
        return f'{self.get_config_path()}/operator_proxy_session.json'

    def get_operator_proxy_session(self):
        try:
            config_file_contents = open(self.get_operator_proxy_session_path(), 'r').read()
        except FileNotFoundError:
            return None
        try:
            data = json.loads(config_file_contents)
        except ValueError:
            return None
        from core.models.OperatorProxySession import OperatorProxySession
        return OperatorProxySession.from_dict(data)

    def has_operator_proxy_session(self):
        return os.path.isfile(self.get_operator_proxy_session_path())

    def __delete_wireguard_configuration(self):
        if self.has_wireguard_configuration():
            if shutil.which('pkexec') is None:
                raise CommandNotFoundError('pkexec')

            try:
                process = subprocess.run(('pkexec', 'rm', '-rf', self.get_wireguard_configuration_path()), check=True)
                completed_successfully = not bool(os.waitpid(process.pid, 0)[1] >> 8)
            except subprocess.CalledProcessError as e:
                completed_successfully = True
            except:
                completed_successfully = True

            if not completed_successfully:
                raise ProfileModificationError('The WireGuard configuration could not be deleted.')

    @staticmethod
    def __get_system_config_path(id: int):
        config_path = f'{Constants.HV_SYSTEM_PROFILE_CONFIG_PATH}/{str(id)}'
        return config_path