src/analyzer/models.py

"""Security analysis data models and logic."""

from dataclasses import dataclass
from enum import Enum
from typing import Optional


class RiskLevel(str, Enum):
    """Severity levels for security incidents."""

    CRITICAL = "critical"
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"
    INFO = "info"


@dataclass
class SecurityAnalysis:
    """Structured analysis result from the LLM."""

    summary: str
    """Brief summary of what happened."""

    risk_level: RiskLevel
    """Severity classification."""

    remediation: str
    """Suggested corrective actions."""

    indicators: list[str]
    """Key indicators of compromise or anomalies found."""

    raw_response: str
    """Full LLM response for transparency."""

    def to_dict(self) -> dict:
        """Convert to dictionary for Gradio output."""
        return {
            "summary": self.summary,
            "risk_level": self.risk_level.value.upper(),
            "remediation": self.remediation,
            "indicators": self.indicators,
        }