backend/database.py

"""
NetGuard Pro — Database Layer
SQLite via SQLAlchemy (async-compatible through run_in_executor).
Tables:
  flows   — every scored flow record
  alerts  — anomalous flows with explanation + severity
  metrics — periodic model performance snapshots
"""
import json
from datetime import datetime, timezone
from sqlalchemy import (
    create_engine, Column, Integer, Float, String,
    Boolean, DateTime, Text, Index
)
from sqlalchemy.orm import declarative_base, sessionmaker

from config import DB_PATH

DATABASE_URL = f"sqlite:///{DB_PATH}"

engine = create_engine(
    DATABASE_URL,
    connect_args={"check_same_thread": False},
    echo=False,
)
SessionLocal = sessionmaker(bind=engine, autocommit=False, autoflush=False)
Base = declarative_base()


def serialize_utc_timestamp(value):
    if not value:
        return None
    if value.tzinfo is None:
        value = value.replace(tzinfo=timezone.utc)
    else:
        value = value.astimezone(timezone.utc)
    return value.isoformat().replace("+00:00", "Z")


# ── ORM Models ──────────────────────────────────────────────────────────────

class FlowRecord(Base):
    __tablename__ = "flows"

    id              = Column(Integer, primary_key=True, index=True)
    flow_id         = Column(String(32), index=True)
    timestamp       = Column(DateTime, default=datetime.utcnow, index=True)

    # Core features (raw, unscaled)
    duration        = Column(Float, default=0)
    packet_rate     = Column(Float, default=0)
    byte_rate       = Column(Float, default=0)
    avg_packet_size = Column(Float, default=0)
    packet_loss     = Column(Float, default=0)
    jitter_ms       = Column(Float, default=0)
    is_tcp          = Column(Boolean, default=False)
    is_udp          = Column(Boolean, default=False)
    flow_asymmetry  = Column(Float, default=0)
    port_entropy    = Column(Float, default=0)
    src_port        = Column(Integer, default=0)
    dst_port        = Column(Integer, default=0)
    src_ip          = Column(String(15), default="")
    dst_ip          = Column(String(15), default="")

    # Prediction
    final_score     = Column(Float, default=0)
    is_anomaly      = Column(Boolean, default=False)
    attack_type     = Column(String(32), default="normal")

    # Per-detector scores (stored as JSON string)
    detector_scores_json = Column(Text, default="{}")

    @property
    def detector_scores(self):
        return json.loads(self.detector_scores_json or "{}")

    @detector_scores.setter
    def detector_scores(self, val):
        self.detector_scores_json = json.dumps(val)

    def to_dict(self):
        return {
            "id":             self.id,
            "flow_id":        self.flow_id,
            "timestamp":      serialize_utc_timestamp(self.timestamp),
            "packet_rate":    self.packet_rate,
            "byte_rate":      self.byte_rate,
            "duration":       self.duration,
            "avg_packet_size":self.avg_packet_size,
            "is_tcp":         self.is_tcp,
            "is_udp":         self.is_udp,
            "src_ip":         self.src_ip,
            "dst_ip":         self.dst_ip,
            "src_port":       self.src_port,
            "dst_port":       self.dst_port,
            "final_score":    self.final_score,
            "is_anomaly":     self.is_anomaly,
            "attack_type":    self.attack_type,
            "detector_scores":self.detector_scores,
        }


class AlertRecord(Base):
    __tablename__ = "alerts"

    id              = Column(Integer, primary_key=True, index=True)
    flow_id         = Column(String(32), index=True)
    timestamp       = Column(DateTime, default=datetime.utcnow, index=True)
    attack_type     = Column(String(32))
    severity        = Column(String(8))          # "high" | "medium" | "low"
    final_score     = Column(Float)
    explanation_json= Column(Text, default="{}")
    raw_flow_json   = Column(Text, default="{}")
    acknowledged    = Column(Boolean, default=False)

    @property
    def explanation(self):
        return json.loads(self.explanation_json or "{}")

    @property
    def raw_flow(self):
        return json.loads(self.raw_flow_json or "{}")

    def to_dict(self):
        return {
            "id":           self.id,
            "flow_id":      self.flow_id,
            "timestamp":    serialize_utc_timestamp(self.timestamp),
            "attack_type":  self.attack_type,
            "severity":     self.severity,
            "final_score":  self.final_score,
            "explanation":  self.explanation,
            "raw_flow":     self.raw_flow,
            "acknowledged": self.acknowledged,
        }


class MetricsSnapshot(Base):
    __tablename__ = "metrics"

    id             = Column(Integer, primary_key=True, index=True)
    timestamp      = Column(DateTime, default=datetime.utcnow, index=True)
    window_flows   = Column(Integer, default=0)
    anomaly_count  = Column(Integer, default=0)
    normal_count   = Column(Integer, default=0)
    avg_score      = Column(Float, default=0)
    threshold      = Column(Float, default=0.5)
    weights_json   = Column(Text, default="{}")

    @property
    def weights(self):
        return json.loads(self.weights_json or "{}")

    def to_dict(self):
        return {
            "timestamp":    serialize_utc_timestamp(self.timestamp),
            "window_flows": self.window_flows,
            "anomaly_count":self.anomaly_count,
            "normal_count": self.normal_count,
            "anomaly_rate": round(self.anomaly_count / max(self.window_flows, 1), 4),
            "avg_score":    round(self.avg_score, 4),
            "threshold":    self.threshold,
            "weights":      self.weights,
        }


# ── Indexes ─────────────────────────────────────────────────────────────────
Index("ix_flows_timestamp_anomaly", FlowRecord.timestamp, FlowRecord.is_anomaly)
Index("ix_alerts_timestamp_ack",    AlertRecord.timestamp, AlertRecord.acknowledged)


def init_db():
    """Create all tables. Called once at startup."""
    Base.metadata.create_all(bind=engine)
    print(f"[DB] Initialized at {DB_PATH}")


def get_db():
    """FastAPI dependency — yields a DB session."""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()