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