admin / Syanpse-Vanguard

"""Normalized security-event model (ECS-flavored) and the raw agent event that
the normalizer turns into it."""
from __future__ import annotations

from datetime import datetime, timezone
from enum import IntEnum
from uuid import UUID

from pydantic import BaseModel, ConfigDict, Field

from sv.config import DEFAULT_TENANT_ID


class Severity(IntEnum):
    EMERGENCY = 0
    ALERT = 1
    CRITICAL = 2
    ERROR = 3
    WARNING = 4
    NOTICE = 5
    INFO = 6
    DEBUG = 7


# Column order is the single source of truth for the Arrow batch and the
# INSERT column list. Keep it in sync with schema.sql.
STORAGE_COLUMNS: list[str] = [
    "ts", "ingested_ts", "tenant_id", "source_type", "host", "host_ip",
    "event_id", "severity", "category", "action", "outcome", "user_name",
    "process_name", "mitre_technique", "message", "raw", "labels",
]


class RawEvent(BaseModel):
    """What an agent ships on the wire — one NDJSON line per event. Deliberately
    permissive; the normalizer is responsible for mapping to EventModel."""
    model_config = ConfigDict(extra="allow")

    source_type: str                     # winlog | journald | syslog | auditd ...
    host: str
    ts: datetime | None = None           # agent-provided event time (optional)
    message: str = ""
    # channel-specific fields land in `fields` or as extras
    fields: dict = Field(default_factory=dict)


class EventModel(BaseModel):
    """The normalized, storable event."""
    model_config = ConfigDict(extra="forbid")

    ts: datetime
    ingested_ts: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    tenant_id: UUID = DEFAULT_TENANT_ID
    source_type: str
    host: str
    host_ip: str | None = None
    event_id: int | None = None
    severity: Severity = Severity.INFO
    category: str | None = None
    action: str | None = None
    outcome: str | None = None
    user_name: str | None = None
    process_name: str | None = None
    mitre_technique: str | None = None
    message: str = ""
    raw: dict = Field(default_factory=dict)
    labels: dict[str, str] = Field(default_factory=dict)

    @staticmethod
    def storage_columns() -> list[str]:
        return STORAGE_COLUMNS

    def to_storage_row(self) -> dict:
        import json
        return {
            "ts": _naive_utc(self.ts),
            "ingested_ts": _naive_utc(self.ingested_ts),
            "tenant_id": str(self.tenant_id),
            "source_type": self.source_type,
            "host": self.host,
            "host_ip": self.host_ip,
            "event_id": self.event_id,
            "severity": int(self.severity),
            "category": self.category,
            "action": self.action,
            "outcome": self.outcome,
            "user_name": self.user_name,
            "process_name": self.process_name,
            "mitre_technique": self.mitre_technique,
            "message": self.message,
            "raw": json.dumps(self.raw, default=str),
            "labels": json.dumps(self.labels, default=str),
        }


def _naive_utc(dt: datetime) -> datetime:
    """DuckDB TIMESTAMP is timezone-naive; store everything as naive UTC."""
    if dt.tzinfo is not None:
        dt = dt.astimezone(timezone.utc).replace(tzinfo=None)
    return dt