VPNTray/models.py

from dataclasses import dataclass, field
from typing import List, Optional
from enum import Enum


class ServiceType(Enum):
    """Enum for different types of services that can run on hosts."""
    SSH = "SSH"
    WEB_GUI = "Web GUI"
    RDP = "RDP"
    VNC = "VNC"
    SMB = "SMB"
    DATABASE = "Database"
    FTP = "FTP"


@dataclass
class Service:
    """Represents a service on a host."""
    name: str
    service_type: ServiceType
    port: int


class HostType(Enum):
    """Enum for different types of hosts."""
    LINUX = "Linux"
    WINDOWS = "Windows"
    WINDOWS_SERVER = "Windows Server"
    PROXMOX = "Proxmox"
    ESXI = "ESXi"
    ROUTER = "Router"
    SWITCH = "Switch"


class VPNType(Enum):
    """Enum for different VPN types."""
    OPENVPN = "OpenVPN"
    WIREGUARD = "WireGuard"
    IPSEC = "IPSec"


@dataclass
class Host:
    """Represents a physical or virtual host at a location."""
    name: str
    ip_address: str
    host_type: HostType
    description: str = ""
    services: List[Service] = field(default_factory=list)
    sub_hosts: List['Host'] = field(
        default_factory=list)  # For VMs under hypervisors

    def get_service_by_name(self, service_name: str) -> Optional[Service]:
        """Get a service by its name."""
        for service in self.services:
            if service.name == service_name:
                return service
        return None

    def is_hypervisor(self) -> bool:
        """Check if this host has sub-hosts (VMs)."""
        return len(self.sub_hosts) > 0


@dataclass
class Location:
    """Represents a customer location."""
    name: str
    vpn_type: VPNType
    connected: bool = False
    active: bool = False
    vpn_config: str = ""  # Path to VPN config or connection details
    hosts: List[Host] = field(default_factory=list)
    
    # VPN connection management fields
    nmcli_connection_name: Optional[str] = None  # NetworkManager connection name
    auto_import: bool = True  # Auto-import .ovpn file if not in NetworkManager
    
    # Credential storage - can be:
    # - Passbolt UUID string (for future use)
    # - Dict with 'username' and 'password' keys
    # - None if no credentials needed
    vpn_credentials: Optional[dict | str] = None

    def get_host_by_name(self, host_name: str) -> Optional[Host]:
        """Get a host by its name (searches recursively in sub-hosts)."""
        def search_hosts(hosts_list: List[Host]) -> Optional[Host]:
            for host in hosts_list:
                if host.name == host_name:
                    return host
                # Search in sub-hosts
                sub_result = search_hosts(host.sub_hosts)
                if sub_result:
                    return sub_result
            return None

        return search_hosts(self.hosts)

    def get_all_hosts_flat(self) -> List[Host]:
        """Get all hosts including sub-hosts in a flat list."""
        def collect_hosts(hosts_list: List[Host]) -> List[Host]:
            result = []
            for host in hosts_list:
                result.append(host)
                result.extend(collect_hosts(host.sub_hosts))
            return result

        return collect_hosts(self.hosts)

    def get_hypervisors(self) -> List[Host]:
        """Get all hosts that have sub-hosts (hypervisors)."""
        return [host for host in self.get_all_hosts_flat() if host.is_hypervisor()]


@dataclass
class CustomerService:
    """Represents a customer's cloud/web service."""
    name: str
    url: str
    service_type: str  # e.g., "Email", "Phone System", "CRM", "ERP"
    description: str = ""


@dataclass
class Customer:
    """Represents a customer with their services and locations."""
    name: str

    # Customer's cloud/web services (available regardless of location)
    services: List[CustomerService] = field(default_factory=list)

    # Customer's locations with their infrastructure
    locations: List[Location] = field(default_factory=list)

    def get_location_by_name(self, location_name: str) -> Optional[Location]:
        """Get a location by its name."""
        for location in self.locations:
            if location.name == location_name:
                return location
        return None

    def get_active_locations(self) -> List[Location]:
        """Get all active locations for this customer."""
        return [loc for loc in self.locations if loc.active]

    def get_inactive_locations(self) -> List[Location]:
        """Get all inactive locations for this customer."""
        return [loc for loc in self.locations if not loc.active]

    def has_active_locations(self) -> bool:
        """Check if customer has any active locations."""
        return any(loc.active for loc in self.locations)

    def has_connected_locations(self) -> bool:
        """Check if customer has any connected locations."""
        return any(loc.connected for loc in self.locations)

    def get_all_hosts_flat(self) -> List[Host]:
        """Get all hosts from all locations in a flat list."""
        all_hosts = []
        for location in self.locations:
            all_hosts.extend(location.get_all_hosts_flat())
        return all_hosts