models.py

"""
Bastion: Cybersecurity Incident Response Environment — Data Models

All models extend OpenEnv base types (Action, Observation, State).
"""

from __future__ import annotations

import random
from enum import IntEnum
from typing import Any, Dict, List, Optional

from pydantic import Field

from openenv.core.env_server import Action, Observation, State


# ---------------------------------------------------------------------------
# Systems in the network
# ---------------------------------------------------------------------------

SYSTEM_NAMES = [
    "web_server",
    "app_server",
    "database",
    "file_server",
    "email_server",
    "workstations",
    "backup_server",
    "firewall",
]
NUM_SYSTEMS = len(SYSTEM_NAMES)

# Which systems can the attacker reach from each system (adjacency)
NETWORK_ADJACENCY: Dict[str, List[str]] = {
    "web_server": ["app_server", "firewall"],
    "app_server": ["web_server", "database", "file_server"],
    "database": ["app_server", "backup_server"],
    "file_server": ["app_server", "email_server", "workstations"],
    "email_server": ["file_server", "workstations"],
    "workstations": ["file_server", "email_server", "app_server"],
    "backup_server": ["database"],
    "firewall": ["web_server"],
}

# How important is each system (affects scoring)
SYSTEM_CRITICALITY: Dict[str, float] = {
    "web_server": 0.6,
    "app_server": 0.8,
    "database": 1.0,
    "file_server": 0.7,
    "email_server": 0.4,
    "workstations": 0.3,
    "backup_server": 0.9,
    "firewall": 0.5,
}

# Systems that hold exfiltrable data
DATA_SYSTEMS = {"database", "file_server", "email_server", "backup_server"}

# Systems that are production-critical services
SERVICE_SYSTEMS = {"web_server", "app_server", "database", "email_server"}


# ---------------------------------------------------------------------------
# Alert model
# ---------------------------------------------------------------------------

class AlertSeverity(IntEnum):
    LOW = 0
    MEDIUM = 1
    HIGH = 2
    CRITICAL = 3


# ---------------------------------------------------------------------------
# Action Space
# ---------------------------------------------------------------------------

class ActionType(IntEnum):
    INVESTIGATE_SYSTEM = 0
    ISOLATE_SYSTEM = 1
    PATCH_VULNERABILITY = 2
    RESTORE_FROM_BACKUP = 3
    ANALYZE_ALERTS = 4
    DEPLOY_MONITORING = 5
    ESCALATE_TO_MANAGEMENT = 6
    BLOCK_EXTERNAL_TRAFFIC = 7
    HUNT_THREAT = 8
    COORDINATE_TEAM = 9


ACTION_NAMES: Dict[int, str] = {int(a): a.name.lower() for a in ActionType}
NUM_ACTIONS = len(ActionType)

# Actions that target a specific system need a target
TARGETED_ACTIONS = {
    ActionType.INVESTIGATE_SYSTEM,
    ActionType.ISOLATE_SYSTEM,
    ActionType.PATCH_VULNERABILITY,
    ActionType.RESTORE_FROM_BACKUP,
    ActionType.HUNT_THREAT,
}


class IncidentAction(Action):
    """Agent picks an action and optionally a target system."""
    action: int = Field(
        ..., ge=0, lt=NUM_ACTIONS,
        description="Action index (0-9). See ActionType enum.",
    )
    target_system: int = Field(
        default=0, ge=0, lt=NUM_SYSTEMS,
        description="Target system index (0-7) for targeted actions. "
                    "0=web_server, 1=app_server, 2=database, 3=file_server, "
                    "4=email_server, 5=workstations, 6=backup_server, 7=firewall",
    )


# ---------------------------------------------------------------------------
# Per-system state
# ---------------------------------------------------------------------------

class SystemState(State):
    """State of a single system in the network."""
    name: str = ""
    compromised: bool = False
    isolated: bool = False
    investigated: bool = False
    has_backdoor: bool = False
    integrity: float = Field(default=1.0, ge=0.0, le=1.0)
    criticality: float = Field(default=0.5, ge=0.0, le=1.0)
    monitoring_level: int = Field(default=0, ge=0, le=3)  # 0=none, 3=full
    patched: bool = False


# ---------------------------------------------------------------------------
# Alert (what the agent sees in the queue)
# ---------------------------------------------------------------------------

class Alert(State):
    """A security alert — may be real or false positive."""
    source_system: str = ""
    severity: int = 0
    message: str = ""
    is_true_positive: bool = True  # hidden from agent in observation
    hour: int = 0
    # SIEM-enriched fields (all optional, backward-compatible)
    source_ip: str = ""
    dest_ip: str = ""
    mitre_technique: str = ""       # e.g. "T1021.002"
    mitre_tactic: str = ""          # e.g. "Lateral Movement"
    process_name: str = ""          # e.g. "powershell.exe"
    event_id: str = ""              # e.g. "EVT-4624"
    file_hash: str = ""             # e.g. "sha256:a1b2c3..."
    confidence: float = 0.0         # detection confidence 0.0-1.0
    raw_log: str = ""               # simulated raw syslog line


# ---------------------------------------------------------------------------
# Full incident state (ground truth)
# ---------------------------------------------------------------------------

class IncidentState(State):
    """Full ground-truth state of the incident response simulation."""
    systems: List[SystemState] = Field(default_factory=list)
    alerts: List[Alert] = Field(default_factory=list)

    # Global metrics
    attacker_progress: float = Field(default=0.0, ge=0.0, le=1.0)
    attacker_stealth: float = Field(default=0.8, ge=0.0, le=1.0)
    data_exfiltrated: float = Field(default=0.0, ge=0.0, le=1.0)
    services_disrupted: int = Field(default=0, ge=0)
    team_stamina: float = Field(default=1.0, ge=0.0, le=1.0)
    hour: int = Field(default=0, ge=0)

    # Tracking
    external_blocked: bool = False
    management_escalated: bool = False
    management_pressure: float = Field(default=0.0, ge=0.0, le=1.0)
    task_id: str = ""

    def get_system(self, name: str) -> SystemState:
        for s in self.systems:
            if s.name == name:
                return s
        raise ValueError(f"Unknown system: {name}")

    def get_system_by_idx(self, idx: int) -> SystemState:
        return self.systems[idx]

    @property
    def compromised_count(self) -> int:
        return sum(1 for s in self.systems if s.compromised and not s.isolated)

    @property
    def isolated_count(self) -> int:
        return sum(1 for s in self.systems if s.isolated)

    @property
    def investigated_count(self) -> int:
        return sum(1 for s in self.systems if s.investigated)

    @property
    def services_intact(self) -> int:
        return sum(
            1 for s in self.systems
            if s.name in SERVICE_SYSTEMS and not s.isolated and s.integrity > 0.3
        )

    def snapshot(self) -> Dict[str, Any]:
        d = self.model_dump()
        d["compromised_count"] = self.compromised_count
        d["isolated_count"] = self.isolated_count
        d["investigated_count"] = self.investigated_count
        d["services_intact"] = self.services_intact
        return d

    def clone(self) -> IncidentState:
        return IncidentState(**self.model_dump())


# ---------------------------------------------------------------------------
# Observation (partial observability)
# ---------------------------------------------------------------------------

class IncidentObservation(Observation):
    """
    What the IR commander sees — noisy, incomplete, delayed.

    Key partial observability:
    - compromise status is UNKNOWN unless system has been investigated
    - data_exfiltrated is estimated with noise
    - alert queue contains false positives
    - attacker_progress is invisible
    """
    # Per-system visible state
    systems_visible: List[Dict[str, Any]] = Field(
        default_factory=list,
        description="Visible state of each system (compromise unknown unless investigated)",
    )
    # Alert queue (includes false positives — agent doesn't know which)
    alert_queue: List[Dict[str, Any]] = Field(
        default_factory=list,
        description="Recent security alerts (may include false positives)",
    )
    # Global estimates (noisy)
    estimated_breach_severity: str = Field(
        default="unknown",
        description="Estimated severity: unknown, low, medium, high, critical",
    )
    estimated_data_at_risk: float = Field(
        default=0.0,
        description="Estimated fraction of data at risk (noisy, 0-1)",
    )
    services_disrupted: int = Field(default=0)
    services_total: int = Field(default=len(SERVICE_SYSTEMS))
    team_stamina: float = Field(default=1.0)
    hour: int = Field(default=0)
    hours_remaining: int = Field(default=12)
    external_blocked: bool = Field(default=False)
    management_escalated: bool = Field(default=False)
    task_description: str = Field(default="")
    # Team communications — messages from IR team members
    team_messages: List[Dict[str, str]] = Field(
        default_factory=list,
        description="Messages from team members (may contain opinions, requests, incorrect assumptions)",
    )


def make_observation(
    state: IncidentState,
    rng: random.Random,
    task_description: str = "",
    done: bool = False,
    reward: float | None = None,
    team_messages: list[dict[str, str]] | None = None,
    alerts_accurate: bool = False,
) -> IncidentObservation:
    """Build a partially-observable view of the true state."""

    # Per-system: only show compromise if investigated
    systems_vis = []
    for s in state.systems:
        vis: Dict[str, Any] = {
            "name": s.name,
            "isolated": s.isolated,
            "investigated": s.investigated,
            "integrity": round(s.integrity, 2),
            "criticality": s.criticality,
            "monitoring_level": s.monitoring_level,
            "patched": s.patched,
        }
        if s.investigated:
            vis["compromised"] = s.compromised
            vis["has_backdoor"] = s.has_backdoor
        else:
            # Unknown — agent must investigate to find out
            vis["compromised"] = "unknown"
            vis["has_backdoor"] = "unknown"
        systems_vis.append(vis)

    # Alerts: show all but hide is_true_positive unless alerts_accurate
    alert_vis = []
    for a in state.alerts[-6:]:  # show last 6 alerts
        av: Dict[str, Any] = {
            "source_system": a.source_system,
            "severity": ["low", "medium", "high", "critical"][a.severity],
            "message": a.message,
            "hour": a.hour,
        }
        # SIEM-enriched fields (only include if populated)
        if a.source_ip:
            av["source_ip"] = a.source_ip
        if a.dest_ip:
            av["dest_ip"] = a.dest_ip
        if a.mitre_technique:
            av["mitre_technique"] = a.mitre_technique
            av["mitre_tactic"] = a.mitre_tactic
        if a.process_name:
            av["process_name"] = a.process_name
        if a.event_id:
            av["event_id"] = a.event_id
        if a.confidence > 0:
            av["confidence"] = round(a.confidence, 2)
        if a.file_hash:
            av["file_hash"] = a.file_hash
        if a.raw_log:
            av["raw_log"] = a.raw_log
        if alerts_accurate:
            av["confirmed"] = a.is_true_positive
        alert_vis.append(av)

    # Estimate breach severity based on what's visible
    investigated = [s for s in state.systems if s.investigated]
    known_compromised = sum(1 for s in investigated if s.compromised)
    if not investigated:
        severity_est = "unknown"
    elif known_compromised == 0:
        severity_est = "low"
    elif known_compromised <= 2:
        severity_est = "medium"
    elif known_compromised <= 4:
        severity_est = "high"
    else:
        severity_est = "critical"

    # Noisy data exfiltration estimate
    noise = rng.gauss(0, 0.1)
    data_est = max(0.0, min(1.0, state.data_exfiltrated + noise))

    # Count disrupted services
    disrupted = sum(
        1 for s in state.systems
        if s.name in SERVICE_SYSTEMS and (s.isolated or s.integrity < 0.3)
    )

    return IncidentObservation(
        systems_visible=systems_vis,
        alert_queue=alert_vis,
        estimated_breach_severity=severity_est,
        estimated_data_at_risk=round(data_est, 3),
        services_disrupted=disrupted,
        team_stamina=round(state.team_stamina, 2),
        hour=state.hour,
        hours_remaining=12 - state.hour,
        external_blocked=state.external_blocked,
        management_escalated=state.management_escalated,
        task_description=task_description,
        team_messages=team_messages or [],
        done=done,
        reward=reward,
    )