Decipher/models.py

from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Any

class Interface(BaseModel):
    name: str
    ip: Optional[str] = None
    subnet: Optional[str] = None
    ipv6: Optional[str] = None
    vlan: Optional[str] = None

class WirelessInfo(BaseModel):
    role: str # "server" or "client"
    ssid: Optional[str] = None
    encryption: Optional[str] = None
    key: Optional[str] = None

class OSPFNetwork(BaseModel):
    network: str

class OSPFProcess(BaseModel):
    process_id: str
    networks: List[str]

class Routing(BaseModel):
    static_routes: List[str] = []
    ospf: List[OSPFProcess] = []
    nat: List[str] = []
    dhcp_pools: List[Dict[str, Any]] = []

class Security(BaseModel):
    acls: List[Dict[str, Any]] = []
    vulnerabilities: List[str] = []
    decrypted_passwords: List[Dict[str, Any]] = []

class Device(BaseModel):
    name: str
    type: str
    model: Optional[str] = ""
    interfaces: List[Interface] = []
    wireless_info: Optional[WirelessInfo] = None
    vlans: List[Dict[str, str]] = []
    routing: Optional[Routing] = None
    security: Optional[Security] = None
    services: List[Dict[str, Any]] = []

class PhysicalLink(BaseModel):
    type: str
    from_device: str
    to_device: str
    from_port: Optional[str] = None
    to_port: Optional[str] = None

class WirelessConnection(BaseModel):
    from_device: str
    to_device: str
    ssid: str
    type: str = "wireless"

class NetworkTopology(BaseModel):
    version: Optional[str] = "1.0"
    devices: List[Device] = []
    physical_links: List[PhysicalLink] = []
    wireless_connections: List[WirelessConnection] = []
    vlans: Dict[str, str] = {}
    instructions: str = ""

    model_config = {
        "json_schema_extra": {
            "example": {
                "version": "1.0",
                "devices": [
                    {
                        "name": "Router-ISP",
                        "type": "Router",
                        "model": "2911",
                        "interfaces": [
                            {"name": "GigabitEthernet0/0", "ip": "10.0.0.1", "subnet": "255.255.255.0"}
                        ]
                    }
                ],
                "physical_links": [],
                "instructions": "This is an example topology."
            }
        }
    }