stuff

models.py

@@ -1,54 +1,154 @@
-from dataclasses import dataclass
+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 host/service within a location."""
+    """Represents a physical or virtual host at a location."""
     name: str
-    address: str
-    type: str  # e.g., "SSH", "Web", "SMB", "PostgreSQL", "Redis"
+    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 location within a customer (e.g., headquarters, branch office)."""
+    """Represents a customer location."""
     name: str
-    vpn_type: str  # e.g., "OpenVPN", "WireGuard", "IPSec"
+    vpn_type: VPNType
     connected: bool = False
     active: bool = False
-    hosts: List[Host] = None
-    
-    def __post_init__(self):
-        if self.hosts is None:
-            self.hosts = []
+    vpn_config: str = ""  # Path to VPN config or connection details
+    hosts: List[Host] = field(default_factory=list)
+
+    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 multiple locations."""
+    """Represents a customer with their services and locations."""
     name: str
-    locations: List[Location] = None
-    
-    def __post_init__(self):
-        if self.locations is None:
-            self.locations = []
-    
-    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)
-    
+
+    # 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
+        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