models.py

"""
Citadel — Data Models (extends Bastion v1 with council + governance + trust + playbook)

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

New vs Bastion v1:
    * Action space grows 10 -> 18 (adds 8 governance actions)
    * IncidentAction carries an optional justification + cited_lessons + gov args
    * CommanderProposal = justified IncidentAction (what Commander sends to Oversight)
    * OversightAction: structured critique (decision + risk_tier + weakness + missing_evidence + counter_proposal)
    * OversightObservation: the meta-view (proposal + justification + policy checks + history)
    * CouncilState: per-episode bookkeeping of proposals/critiques/revisions
    * IncidentState gains governance_state, trust_state, council_state, adversary_gen
    * IncidentObservation gains shared_playbook, trust_state, governance_summary, adversary_gen
"""

from __future__ import annotations

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

from pydantic import BaseModel, ConfigDict, Field

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

from governance import (
    GovernanceState,
    NUM_GOVERNANCE_ACTIONS,
    FIRST_GOVERNANCE_ACTION,
    TOTAL_ACTIONS,
)
from trust import TrustState
from stakeholder_events import StakeholderState
from investor_agent import InvestorState


# ---------------------------------------------------------------------------
# Systems in the network (unchanged from Bastion v1)
# ---------------------------------------------------------------------------

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

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"],
}

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,
}

DATA_SYSTEMS = {"database", "file_server", "email_server", "backup_server"}
SERVICE_SYSTEMS = {"web_server", "app_server", "database", "email_server"}


# ---------------------------------------------------------------------------
# Alert severity
# ---------------------------------------------------------------------------

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


# ---------------------------------------------------------------------------
# Action Space — 18 actions (10 Bastion + 8 governance)
# ---------------------------------------------------------------------------

class ActionType(IntEnum):
    # Bastion v1 actions (0-9)
    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
    # Governance (10-17) — mirror GovernanceAction enum values
    OPEN_SERVICENOW_INCIDENT = 10
    REQUEST_CAB_APPROVAL = 11
    NOTIFY_SLACK_CHANNEL = 12
    LOG_TO_SOX_AUDIT = 13
    PAGE_ONCALL = 14
    NOTIFY_DATA_OWNER = 15
    START_LEGAL_HOLD = 16
    CLOSE_INCIDENT = 17


ACTION_NAMES: Dict[int, str] = {int(a): a.name.lower() for a in ActionType}
NUM_ACTIONS = len(ActionType)           # 18
BASTION_V1_ACTIONS = 10                 # 0-9 inclusive

assert NUM_ACTIONS == TOTAL_ACTIONS, (
    f"ActionType count {NUM_ACTIONS} must match governance TOTAL_ACTIONS {TOTAL_ACTIONS}"
)

TARGETED_ACTIONS = {
    ActionType.INVESTIGATE_SYSTEM,
    ActionType.ISOLATE_SYSTEM,
    ActionType.PATCH_VULNERABILITY,
    ActionType.RESTORE_FROM_BACKUP,
    ActionType.HUNT_THREAT,
    ActionType.NOTIFY_DATA_OWNER,
    ActionType.REQUEST_CAB_APPROVAL,
}


def is_bastion_action(action_idx: int) -> bool:
    return 0 <= action_idx < BASTION_V1_ACTIONS


def is_governance_action(action_idx: int) -> bool:
    return BASTION_V1_ACTIONS <= action_idx < NUM_ACTIONS


# ---------------------------------------------------------------------------
# IncidentAction (extended) — what Commander submits to the env
# ---------------------------------------------------------------------------

class IsolationMethod(str):
    """
    How to isolate a system. Each method has different speed/reversibility tradeoffs.

      firewall_acl  — block at network layer; reversible, minimal service disruption,
                      but attacker retains process-level access until cleaned
      network_unplug — physical/vlan removal; fastest containment, full service loss,
                       requires on-site to restore (costs an extra step)
      process_kill  — terminate attacker processes only; keeps service up, but risks
                      missing persistence (backdoors survive); 40% chance attacker
                      re-establishes within 2 hours
    """
    FIREWALL_ACL   = "firewall_acl"
    NETWORK_UNPLUG = "network_unplug"
    PROCESS_KILL   = "process_kill"

ISOLATION_METHODS = {IsolationMethod.FIREWALL_ACL, IsolationMethod.NETWORK_UNPLUG, IsolationMethod.PROCESS_KILL}


class PatchStrategy(str):
    """
    How to patch. Affects success probability and service impact.

      hotpatch   — live patching without restart; low disruption, 60% effectiveness
      cold_patch — full restart with patch; high disruption, 90% effectiveness
      virtual_patch — WAF/IDS rule blocks exploit path without touching binary;
                      no disruption, 75% effectiveness, only works for network-facing systems
    """
    HOTPATCH       = "hotpatch"
    COLD_PATCH     = "cold_patch"
    VIRTUAL_PATCH  = "virtual_patch"

PATCH_STRATEGIES = {PatchStrategy.HOTPATCH, PatchStrategy.COLD_PATCH, PatchStrategy.VIRTUAL_PATCH}


class MonitoringScope(str):
    """
    What to monitor. Affects detection quality vs noise ratio.

      process_events   — log all process spawns; catches lateral movement, high volume
      network_traffic  — capture all flows; catches exfil, very high volume
      auth_events      — log all auth attempts; targeted, low volume, misses post-auth
      full_endpoint    — all of the above; maximum detection, significant performance hit
                         (integrity -5% per hour on target while active)
    """
    PROCESS_EVENTS  = "process_events"
    NETWORK_TRAFFIC = "network_traffic"
    AUTH_EVENTS     = "auth_events"
    FULL_ENDPOINT   = "full_endpoint"

MONITORING_SCOPES = {MonitoringScope.PROCESS_EVENTS, MonitoringScope.NETWORK_TRAFFIC, MonitoringScope.AUTH_EVENTS, MonitoringScope.FULL_ENDPOINT}


class IncidentAction(Action):
    """
    Commander's action. Carries optional justification + cited lessons
    + governance-specific arguments + rich method/scope/rollback fields
    that give the action real technical depth.

    New fields (Option A — richer payloads):
      method       — HOW to execute the action (isolation method, patch strategy, monitoring scope)
      scope        — WHAT scope: IP ranges, process names, port filters
      rollback_plan — HOW to undo if it goes wrong
    """
    action: int = Field(
        ..., ge=0, lt=NUM_ACTIONS,
        description=f"Action index (0-{NUM_ACTIONS - 1}). 0-9 = incident response, 10-17 = governance.",
    )
    target_system: int = Field(
        default=0, ge=0, lt=NUM_SYSTEMS,
        description="Target system index (0-7) for targeted actions.",
    )
    # Council fields — optional, ignored by the Bastion-style server flow
    justification: str = Field(
        default="",
        description="Free-text reasoning for this action (required when council is active).",
        max_length=1000,
    )
    cited_lessons: List[str] = Field(
        default_factory=list,
        description="Lesson IDs from the shared playbook that informed this proposal.",
    )
    # Richer action payload — Option A
    method: str = Field(
        default="",
        max_length=32,
        description=(
            "How to execute the action. "
            "isolate: firewall_acl | network_unplug | process_kill. "
            "patch: hotpatch | cold_patch | virtual_patch. "
            "deploy_monitoring: process_events | network_traffic | auth_events | full_endpoint."
        ),
    )
    scope: str = Field(
        default="",
        max_length=200,
        description="Optional scope constraint: IP range, process name filter, port. E.g. '10.1.3.30/32' or 'lsass.exe'.",
    )
    rollback_plan: str = Field(
        default="",
        max_length=300,
        description="How to reverse this action if it causes collateral damage or was wrong.",
    )
    # Governance-specific args — used by actions 10-17, ignored otherwise
    severity_arg: int = Field(default=2, ge=1, le=4, description="For governance actions (P1-P4).")
    channel_arg: str = Field(default="sec-ops", max_length=64)
    message_arg: str = Field(default="", max_length=400)
    scope_arg: str = Field(default="", max_length=200)
    evidence_arg: str = Field(default="", max_length=400)


# ---------------------------------------------------------------------------
# CommanderProposal — the wrapper the env uses to talk to Oversight
# ---------------------------------------------------------------------------

class CommanderProposal(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)
    """What the Commander proposes in step t. Oversight sees one of these."""
    action: int
    target_system: int
    justification: str = ""
    cited_lessons: List[str] = Field(default_factory=list)
    method: str = ""
    scope: str = ""
    rollback_plan: str = ""
    severity_arg: int = 2
    channel_arg: str = "sec-ops"
    message_arg: str = ""
    scope_arg: str = ""
    evidence_arg: str = ""

    @classmethod
    def from_action(cls, a: IncidentAction) -> "CommanderProposal":
        return cls(
            action=a.action,
            target_system=a.target_system,
            justification=a.justification,
            cited_lessons=list(a.cited_lessons),
            method=a.method,
            scope=a.scope,
            rollback_plan=a.rollback_plan,
            severity_arg=a.severity_arg,
            channel_arg=a.channel_arg,
            message_arg=a.message_arg,
            scope_arg=a.scope_arg,
            evidence_arg=a.evidence_arg,
        )

    def to_action(self) -> IncidentAction:
        return IncidentAction(
            action=self.action,
            target_system=self.target_system,
            justification=self.justification,
            cited_lessons=list(self.cited_lessons),
            method=self.method,
            scope=self.scope,
            rollback_plan=self.rollback_plan,
            severity_arg=self.severity_arg,
            channel_arg=self.channel_arg,
            message_arg=self.message_arg,
            scope_arg=self.scope_arg,
            evidence_arg=self.evidence_arg,
        )


# ---------------------------------------------------------------------------
# Oversight action space — council decision + structured critique
# ---------------------------------------------------------------------------

class OversightDecision(IntEnum):
    APPROVE = 0
    REVISE = 1               # ask Commander to revise once before deciding
    VETO = 2                 # hard reject — Commander must pick a DIFFERENT action
    FLAG_FOR_HUMAN = 3       # execute but audit-flag

NUM_OVERSIGHT_DECISIONS = len(OversightDecision)


class CounterProposal(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)
    """Oversight's suggested alternative (used with REVISE or VETO)."""
    action: int = Field(..., ge=0, lt=NUM_ACTIONS)
    target_system: int = Field(default=0, ge=0, lt=NUM_SYSTEMS)
    rationale: str = Field(default="", max_length=400)


class OversightAction(Action):
    """
    The structured critique Oversight emits every step. This is the upgrade
    over a classifier-style gate: Oversight must *reason about the domain*.
    """
    decision: int = Field(..., ge=0, lt=NUM_OVERSIGHT_DECISIONS,
                          description="0=approve, 1=revise, 2=veto, 3=flag_for_human")
    risk_tier: int = Field(default=2, ge=1, le=5,
                           description="Commander's proposal risk (1=safe, 5=catastrophic)")
    weakness: str = Field(default="", max_length=400,
                          description="Specific weakness in the proposal — be concrete.")
    missing_evidence: List[str] = Field(
        default_factory=list,
        description="Evidence (alerts/logs/investigations) missing to support this action.",
    )
    counter_proposal: Optional[CounterProposal] = Field(
        default=None,
        description="Optional alternative action (required on VETO, optional on REVISE).",
    )
    # Post-mortem lesson — filled in after action resolves (next step).
    lesson_text: str = Field(default="", max_length=240)
    lesson_tags: List[str] = Field(default_factory=list)


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

class SystemState(State):
    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)
    patched: bool = False


# ---------------------------------------------------------------------------
# Alert
# ---------------------------------------------------------------------------

class Alert(State):
    source_system: str = ""
    severity: int = 0
    message: str = ""
    is_true_positive: bool = True
    hour: int = 0
    source_ip: str = ""
    dest_ip: str = ""
    mitre_technique: str = ""
    mitre_tactic: str = ""
    process_name: str = ""
    event_id: str = ""
    file_hash: str = ""
    confidence: float = 0.0
    raw_log: str = ""


# ---------------------------------------------------------------------------
# Council state — per-episode bookkeeping of the Commander↔Oversight loop
# ---------------------------------------------------------------------------

class ProposalRecord(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)
    """One round of the council loop."""
    step: int
    proposal: CommanderProposal
    oversight_decision: int = -1
    oversight_risk_tier: int = 0
    oversight_weakness: str = ""
    oversight_counter_action: int = -1
    revised: bool = False
    final_action: int = -1
    final_target: int = -1
    outcome_correct: Optional[bool] = None
    cited_lessons: List[str] = Field(default_factory=list)


class CouncilState(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)
    """Running record of the council's decisions this episode."""
    history: List[ProposalRecord] = Field(default_factory=list)
    total_proposals: int = 0
    approvals: int = 0
    revisions: int = 0
    vetoes: int = 0
    flags: int = 0
    correct_vetoes: int = 0
    false_vetoes: int = 0
    useful_revisions: int = 0
    vague_critiques: int = 0
    counter_proposals_adopted: int = 0
    counter_proposals_succeeded: int = 0
    lessons_cited: int = 0
    lessons_cited_and_helpful: int = 0

    def log(self, record: ProposalRecord) -> None:
        self.history.append(record)
        self.total_proposals += 1


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

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

    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)

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

    # --- Citadel additions -----------------------------------------------
    governance_state: GovernanceState = Field(default_factory=GovernanceState)
    trust_state: TrustState = Field(default_factory=TrustState)
    council_state: CouncilState = Field(default_factory=CouncilState)
    stakeholder_state: StakeholderState = Field(default_factory=StakeholderState)
    investor_state: InvestorState = Field(default_factory=InvestorState)
    adversary_gen: int = Field(default=1, ge=1, le=4,
        description="Adversary generation this episode (1/2/3 scripted, 4 LLM).")
    episode_id: str = ""
    step_count: int = 0

    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())


# ---------------------------------------------------------------------------
# Observations — Commander + Oversight
# ---------------------------------------------------------------------------

class IncidentObservation(Observation):
    """What the Commander sees at step t. Extended with council surfaces."""
    # Bastion v1 surfaces
    systems_visible: List[Dict[str, Any]] = Field(default_factory=list)
    alert_queue: List[Dict[str, Any]] = Field(default_factory=list)
    estimated_breach_severity: str = "unknown"
    estimated_data_at_risk: float = 0.0
    services_disrupted: int = 0
    services_total: int = Field(default=len(SERVICE_SYSTEMS))
    team_stamina: float = 1.0
    hour: int = 0
    hours_remaining: int = 12
    external_blocked: bool = False
    management_escalated: bool = False
    task_description: str = ""
    team_messages: List[Dict[str, str]] = Field(default_factory=list)

    # --- Citadel additions -----------------------------------------------
    governance_summary: Dict[str, Any] = Field(default_factory=dict)
    trust_summary: Dict[str, Any] = Field(default_factory=dict)
    stakeholder_summary: Dict[str, Any] = Field(default_factory=dict)
    investor_summary: Dict[str, Any] = Field(default_factory=dict)
    shared_playbook: List[Dict[str, Any]] = Field(default_factory=list)
    adversary_gen: int = 1
    last_oversight_critique: Dict[str, Any] = Field(default_factory=dict)
    done: bool = False
    reward: Optional[float] = None
    oversight_reward: Optional[float] = None
    metadata: Optional[Dict[str, Any]] = None


class OversightObservation(Observation):
    """The meta-view — what Oversight sees when asked to critique a proposal."""
    proposed_action: Dict[str, Any] = Field(default_factory=dict)
    justification: str = ""
    cited_lessons: List[str] = Field(default_factory=list)
    commander_observation: Dict[str, Any] = Field(default_factory=dict)
    commander_action_history: List[Dict[str, Any]] = Field(default_factory=list)
    policy_checks: Dict[str, Any] = Field(default_factory=dict)
    veto_budget_remaining: int = 0
    flag_budget_remaining: int = 0
    shared_playbook: List[Dict[str, Any]] = Field(default_factory=list)
    trust_summary: Dict[str, Any] = Field(default_factory=dict)
    governance_summary: Dict[str, Any] = Field(default_factory=dict)
    adversary_gen: int = 1
    hour: int = 0
    task_description: str = ""


# ---------------------------------------------------------------------------
# make_observation — Commander's partially-observable view
# ---------------------------------------------------------------------------

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,
    shared_playbook: list[dict[str, Any]] | None = None,
    last_oversight_critique: dict[str, Any] | None = None,
    hide_adversary_gen: bool = False,
) -> IncidentObservation:
    """Build the Commander's partially-observable view of the true state."""

    systems_vis: List[Dict[str, Any]] = []
    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:
            vis["compromised"] = "unknown"
            vis["has_backdoor"] = "unknown"
        systems_vis.append(vis)

    alert_vis: List[Dict[str, Any]] = []
    for a in state.alerts[-6:]:
        av: Dict[str, Any] = {
            "source_system": a.source_system,
            "severity": ["low", "medium", "high", "critical"][a.severity],
            "message": a.message,
            "hour": a.hour,
        }
        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)

    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"

    noise = rng.gauss(0, 0.1)
    data_est = max(0.0, min(1.0, state.data_exfiltrated + noise))

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

    state.governance_state.set_current_hour(state.hour)
    governance_summary = state.governance_state.snapshot()
    trust_summary = state.trust_state.as_dict()
    stakeholder_summary = state.stakeholder_state.snapshot()
    investor_summary = state.investor_state.snapshot()

    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=max(0, 12 - state.hour),
        external_blocked=state.external_blocked,
        management_escalated=state.management_escalated,
        task_description=task_description,
        team_messages=team_messages or [],
        governance_summary=governance_summary,
        trust_summary=trust_summary,
        stakeholder_summary=stakeholder_summary,
        investor_summary=investor_summary,
        shared_playbook=shared_playbook or [],
        adversary_gen=(0 if hide_adversary_gen else state.adversary_gen),
        last_oversight_critique=last_oversight_critique or {},
        done=done,
        reward=reward,
    )


# ---------------------------------------------------------------------------
# Re-export Lesson from playbook for external imports
# ---------------------------------------------------------------------------

from playbook import Lesson  # noqa: E402,F401  (intentional re-export)