netviz/app/models/network.py

"""
Network models for topology visualization.
"""
from datetime import datetime
from typing import Dict, List, Any, Optional
import json

from app.extensions import db


class Network(db.Model):
    """Network model representing a network topology."""
    __tablename__ = "networks"
    
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(128), nullable=False)
    description = db.Column(db.Text)
    created_at = db.Column(db.DateTime, default=datetime.utcnow, nullable=False)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow, nullable=False)
    
    # Ownership
    user_id = db.Column(db.Integer, db.ForeignKey("users.id"), nullable=False)
    owner = db.relationship("User", back_populates="networks")
    
    # Related objects
    subnets = db.relationship("Subnet", back_populates="network", cascade="all, delete-orphan")
    devices = db.relationship("Device", back_populates="network", cascade="all, delete-orphan")
    firewall_rules = db.relationship("FirewallRule", back_populates="network", cascade="all, delete-orphan")
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert network to dictionary for API responses and exports."""
        return {
            "id": self.id,
            "name": self.name,
            "description": self.description,
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
            "subnets": [subnet.to_dict() for subnet in self.subnets],
            "devices": [device.to_dict() for device in self.devices],
            "firewall_rules": [rule.to_dict() for rule in self.firewall_rules]
        }
    
    def to_json(self) -> str:
        """Convert network to JSON string for exports."""
        return json.dumps(self.to_dict(), indent=2)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any], user_id: int) -> "Network":
        """Create a network from a dictionary (for imports)."""
        network = cls(
            name=data["name"],
            description=data.get("description", ""),
            user_id=user_id
        )
        
        return network

    def __repr__(self) -> str:
        return f"<Network {self.name}>"


class Subnet(db.Model):
    """Subnet model representing a network subnet."""
    __tablename__ = "subnets"
    
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(128), nullable=False)
    cidr = db.Column(db.String(64), nullable=False)
    vlan = db.Column(db.Integer)
    description = db.Column(db.Text)
    
    # Parent network
    network_id = db.Column(db.Integer, db.ForeignKey("networks.id"), nullable=False)
    network = db.relationship("Network", back_populates="subnets")
    
    # Devices in this subnet
    devices = db.relationship("Device", back_populates="subnet")
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert subnet to dictionary."""
        return {
            "id": self.id,
            "name": self.name,
            "cidr": self.cidr,
            "vlan": self.vlan,
            "description": self.description
        }
    
    def __repr__(self) -> str:
        return f"<Subnet {self.cidr}>"


class Device(db.Model):
    """Device model representing a network device or host."""
    __tablename__ = "devices"
    
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(128), nullable=False)
    ip_address = db.Column(db.String(64))
    mac_address = db.Column(db.String(64))
    device_type = db.Column(db.String(64))  # server, router, switch, etc.
    os = db.Column(db.String(128))
    description = db.Column(db.Text)
    
    # Parent network
    network_id = db.Column(db.Integer, db.ForeignKey("networks.id"), nullable=False)
    network = db.relationship("Network", back_populates="devices")
    
    # Subnet membership
    subnet_id = db.Column(db.Integer, db.ForeignKey("subnets.id"))
    subnet = db.relationship("Subnet", back_populates="devices")
    
    # Additional properties stored as JSON
    properties = db.Column(db.Text)
    
    def get_properties(self) -> Dict[str, Any]:
        """Get device properties from JSON field."""
        if not self.properties:
            return {}
        return json.loads(self.properties)
    
    def set_properties(self, properties: Dict[str, Any]) -> None:
        """Set device properties as JSON."""
        self.properties = json.dumps(properties)
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert device to dictionary."""
        return {
            "id": self.id,
            "name": self.name,
            "ip_address": self.ip_address,
            "mac_address": self.mac_address,
            "device_type": self.device_type,
            "os": self.os,
            "description": self.description,
            "subnet_id": self.subnet_id,
            "properties": self.get_properties()
        }
    
    def __repr__(self) -> str:
        return f"<Device {self.name} ({self.ip_address})>"


class FirewallRule(db.Model):
    """Firewall rule model for documenting network security policies."""
    __tablename__ = "firewall_rules"
    
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(128), nullable=False)
    source = db.Column(db.String(128), nullable=False)
    destination = db.Column(db.String(128), nullable=False)
    protocol = db.Column(db.String(16))  # tcp, udp, icmp, any
    port_range = db.Column(db.String(64))  # e.g., "80,443" or "1024-2048"
    action = db.Column(db.String(16), nullable=False)  # allow, deny
    description = db.Column(db.Text)
    
    # Parent network
    network_id = db.Column(db.Integer, db.ForeignKey("networks.id"), nullable=False)
    network = db.relationship("Network", back_populates="firewall_rules")
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert firewall rule to dictionary."""
        return {
            "id": self.id,
            "name": self.name,
            "source": self.source,
            "destination": self.destination,
            "protocol": self.protocol,
            "port_range": self.port_range,
            "action": self.action,
            "description": self.description
        }
    
    def __repr__(self) -> str:
        return f"<FirewallRule {self.name}: {self.source} to {self.destination}>"