soc_env/environment.py

"""
SOCEnv — main OpenEnv environment class.

Public interface (OpenEnv spec):
  reset()        -> Observation
  step(action)   -> (Observation, Reward, done: bool, info: dict)
  state()        -> EnvState
  grade()        -> float  [0.0, 1.0]
"""
from __future__ import annotations
import copy
import json
import os
import sys
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Tuple

from .models import (
    Action, ActionType, AlertStatus, Observation, Reward, EnvState,
    InvestigationNote, MITRETactic, HostStatus, BusinessConstraint,
    NetworkHost, SIEMAlert,
)


def _ts() -> str:
    return datetime.now(timezone.utc).isoformat()


def _load_scenario(task_id: str, seed: int) -> Dict[str, Any]:
    """Import scenarios module with robust path resolution."""
    # Add repo root to path so 'scenarios' package is always found
    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    if root not in sys.path:
        sys.path.insert(0, root)
    from scenarios import load_scenario
    return load_scenario(task_id, seed=seed)


class SOCEnv:
    """Security Operations Center OpenEnv environment."""

    TASK_IDS = [
        "alert_triage",
        "attack_chain_reconstruction",
        "constrained_incident_response",
    ]
    MAX_STEPS = {
        "alert_triage": 10,
        "attack_chain_reconstruction": 25,
        "constrained_incident_response": 40,
    }
    MAX_MINUTES = {
        "alert_triage": 30,
        "attack_chain_reconstruction": 120,
        "constrained_incident_response": 240,
    }
    MINUTES_PER_STEP = {
        "alert_triage": 3,
        "attack_chain_reconstruction": 5,
        "constrained_incident_response": 6,
    }

    def __init__(self, task_id: str = "alert_triage", seed: int = 42):
        if task_id not in self.TASK_IDS:
            raise ValueError(f"task_id must be one of {self.TASK_IDS}")
        self.task_id = task_id
        self.seed = seed
        self._state: Optional[EnvState] = None

    # ------------------------------------------------------------------
    # OpenEnv public interface
    # ------------------------------------------------------------------

    def reset(self) -> Observation:
        """Initialise episode. Returns first Observation."""
        scenario = _load_scenario(self.task_id, self.seed)
        self._state = EnvState(
            task_id=self.task_id,
            step=0,
            done=False,
            cumulative_reward=0.0,
            observation=scenario["observation"],
            attack_chain=scenario.get("attack_chain"),
            agent_classifications={},
            identified_stages=[],
            isolated_hosts=[],
            disabled_accounts=[],
            forensics_collected={},
            escalated=False,
            ticket_created=False,
        )
        return copy.deepcopy(self._state.observation)

    def step(self, action: Action) -> Tuple[Observation, Reward, bool, Dict[str, Any]]:
        """Advance environment one step. Returns (obs, reward, done, info)."""
        if self._state is None:
            raise RuntimeError("Call reset() before step().")
        if self._state.done:
            raise RuntimeError("Episode done. Call reset() to start a new episode.")

        s = self._state
        s.step += 1
        s.observation.step = s.step
        s.observation.steps_remaining = self.MAX_STEPS[self.task_id] - s.step
        s.observation.elapsed_minutes = min(
            s.observation.elapsed_minutes + self.MINUTES_PER_STEP[self.task_id],
            self.MAX_MINUTES[self.task_id],
        )

        action_result, action_ok = self._apply_action(action, s)
        s.observation.last_action_result = action_result
        s.observation.last_action_success = action_ok

        s.observation.notes.append(InvestigationNote(
            step=s.step,
            action_taken=action.action_type.value,
            finding=action_result,
            timestamp=_ts(),
        ))

        reward = self._compute_reward(action, s, action_ok)
        s.cumulative_reward = round(s.cumulative_reward + reward.total, 4)
        s.done = self._check_done(s)

        # Move non-pending alerts to acknowledged list
        s.observation.acknowledged_alerts = [
            a for a in s.observation.active_alerts + s.observation.acknowledged_alerts
            if a.status != AlertStatus.PENDING
        ]
        s.observation.active_alerts = [
            a for a in s.observation.active_alerts
            if a.status == AlertStatus.PENDING
        ]

        info = {
            "cumulative_reward": s.cumulative_reward,
            "step": s.step,
            "identified_stages": [st.value for st in s.identified_stages],
            "isolated_hosts": s.isolated_hosts,
            "done": s.done,
        }
        return copy.deepcopy(s.observation), reward, s.done, info

    def state(self) -> EnvState:
        """Return full environment state including ground truth (for grading)."""
        if self._state is None:
            raise RuntimeError("Call reset() before state().")
        return copy.deepcopy(self._state)

    def grade(self) -> float:
        """Run the grader on current state. Returns float in [0.0, 1.0]."""
        if self._state is None:
            raise RuntimeError("Call reset() and run an episode before grading.")
        from .graders import grade_task_easy, grade_task_medium, grade_task_hard
        s = self._state
        if self.task_id == "alert_triage":
            return grade_task_easy(s)
        elif self.task_id == "attack_chain_reconstruction":
            return grade_task_medium(s)
        else:
            return grade_task_hard(s)

    # ------------------------------------------------------------------
    # Action handlers
    # ------------------------------------------------------------------

    def _apply_action(self, action: Action, s: EnvState) -> Tuple[str, bool]:
        t = action.action_type
        if t == ActionType.ENRICH_ALERT:
            return self._do_enrich(action, s)
        elif t == ActionType.CORRELATE_ALERTS:
            return self._do_correlate(action, s)
        elif t == ActionType.ISOLATE_ENDPOINT:
            return self._do_isolate(action, s)
        elif t == ActionType.DISABLE_ACCOUNT:
            return self._do_disable_account(action, s)
        elif t == ActionType.COLLECT_FORENSICS:
            return self._do_forensics(action, s)
        elif t == ActionType.ESCALATE_TO_TIER2:
            return self._do_escalate(action, s)
        elif t == ActionType.CREATE_TICKET:
            return self._do_ticket(action, s)
        return "Unknown action type.", False

    def _do_enrich(self, action: Action, s: EnvState) -> Tuple[str, bool]:
        alert = self._find_alert(action.alert_id, s)
        if not alert:
            return f"Alert {action.alert_id} not found.", False
        alert.status = AlertStatus.INVESTIGATING
        source = action.source or "threat_intel"
        enrichment: Dict[str, Any] = {"source": source, "retrieved_at": _ts()}
        if source == "threat_intel":
            enrichment["indicators_checked"] = len(alert.indicators)
            enrichment["malicious_hits"] = sum(
                1 for i in alert.indicators if i.reputation == "malicious"
            )
            for ind in alert.indicators:
                ind.reputation = ind.reputation or "unknown"
        elif source == "user_context":
            enrichment["user_id"] = alert.user_id
            enrichment["anomalous_behaviour"] = bool(alert.ground_truth)
        elif source == "asset_db":
            host = self._find_host(alert.host_id, s)
            enrichment["host_role"] = host.role if host else "unknown"
            enrichment["is_critical"] = host.is_critical if host else False
        alert.enrichment = enrichment
        return f"Enriched alert {action.alert_id} via {source}: {json.dumps(enrichment)}", True

    def _do_correlate(self, action: Action, s: EnvState) -> Tuple[str, bool]:
        if not action.alert_ids or len(action.alert_ids) < 2:
            return "correlate_alerts requires at least 2 alert_ids.", False
        all_alerts = s.observation.active_alerts + s.observation.acknowledged_alerts
        alerts = [a for a in all_alerts if a.alert_id in action.alert_ids]
        if len(alerts) < 2:
            return "Could not find 2+ valid alerts to correlate.", False
        new_tactics: List[MITRETactic] = []
        for a in alerts:
            if a.mitre_tactic and a.mitre_tactic not in s.identified_stages:
                s.identified_stages.append(a.mitre_tactic)
                new_tactics.append(a.mitre_tactic)
        return (
            f"Correlated {len(alerts)} alerts. "
            f"Hypothesis: {action.correlation_hypothesis or '(none)'}. "
            f"New ATT&CK tactics: {[t.value for t in new_tactics] or 'none'}."
        ), True

    def _do_isolate(self, action: Action, s: EnvState) -> Tuple[str, bool]:
        host = self._find_host(action.host_id, s)
        if not host:
            return f"Host {action.host_id} not found.", False
        for bc in s.observation.business_constraints:
            if bc.host_id == action.host_id and bc.constraint_type == "cannot_isolate":
                if bc.severity == "hard_block":
                    return f"BLOCKED: Cannot isolate {host.hostname}. Reason: {bc.reason}", False
        if host.status == HostStatus.ISOLATED:
            return f"{host.hostname} is already isolated.", False
        host.status = HostStatus.ISOLATED
        if action.host_id not in s.isolated_hosts:
            s.isolated_hosts.append(action.host_id)
        return f"Endpoint {host.hostname} ({host.ip_address}) isolated from network.", True

    def _do_disable_account(self, action: Action, s: EnvState) -> Tuple[str, bool]:
        if not action.user_id:
            return "user_id is required for disable_account.", False
        if action.user_id not in s.disabled_accounts:
            s.disabled_accounts.append(action.user_id)
        return f"Account {action.user_id} disabled. Active sessions terminated.", True

    def _do_forensics(self, action: Action, s: EnvState) -> Tuple[str, bool]:
        host = self._find_host(action.host_id, s)
        if not host:
            return f"Host {action.host_id} not found.", False
        artifacts = action.artifact_types or ["memory_dump", "event_logs"]
        s.forensics_collected[action.host_id] = list(
            set(s.forensics_collected.get(action.host_id, []) + artifacts)
        )
        for bc in s.observation.business_constraints:
            if bc.host_id == action.host_id and bc.constraint_type == "legal_hold":
                return (
                    f"Forensics collected from {host.hostname}: {artifacts}. "
                    f"LEGAL HOLD applied — evidence preserved per: {bc.reason}."
                ), True
        return f"Forensics collected from {host.hostname}: {artifacts}.", True

    def _do_escalate(self, action: Action, s: EnvState) -> Tuple[str, bool]:
        if s.escalated:
            return "Already escalated to Tier 2.", False
        s.escalated = True
        return f"Escalated to Tier 2. Summary: {action.summary or '(no summary provided)'}.", True

    def _do_ticket(self, action: Action, s: EnvState) -> Tuple[str, bool]:
        s.ticket_created = True
        return (
            f"Ticket created — priority {action.priority or 'P2'}. "
            f"Summary: {action.summary or '(no summary)'}."
        ), True

    # ------------------------------------------------------------------
    # Reward computation
    # ------------------------------------------------------------------

    def _compute_reward(self, action: Action, s: EnvState, action_ok: bool) -> Reward:
        r = Reward(total=0.0)
        if not action_ok:
            r.total = -0.02
            r.info = "invalid_action"
            return r

        t = action.action_type

        if t == ActionType.ENRICH_ALERT:
            r.classification_delta = 0.05
            r.info = "enrichment"

        elif t == ActionType.CORRELATE_ALERTS:
            r.chain_coverage_delta = round(len(s.identified_stages) * 0.08, 3)
            r.info = f"correlation_stages_{len(s.identified_stages)}"

        elif t == ActionType.ISOLATE_ENDPOINT:
            host = self._find_host(action.host_id, s)
            if host:
                chain = s.attack_chain
                in_chain = chain and (
                    action.host_id == chain.patient_zero_host
                    or action.host_id in chain.lateral_movement_targets
                )
                if in_chain:
                    r.containment_delta = 0.25
                    r.info = "correct_isolation"
                else:
                    r.fp_penalty = -0.20
                    r.info = "fp_isolation"

        elif t == ActionType.DISABLE_ACCOUNT:
            all_alerts = s.observation.active_alerts + s.observation.acknowledged_alerts
            alert_users = {a.user_id for a in all_alerts if a.user_id}
            if action.user_id in alert_users:
                r.containment_delta = 0.10
                r.info = "correct_account_disable"
            else:
                r.fp_penalty = -0.05
                r.info = "fp_account_disable"

        elif t == ActionType.COLLECT_FORENSICS:
            r.classification_delta = 0.05
            r.info = "forensics_collected"

        elif t == ActionType.ESCALATE_TO_TIER2:
            if s.observation.business_constraints:
                r.compliance_delta = 0.10
                r.info = "appropriate_escalation"
            else:
                r.fp_penalty = -0.05
                r.info = "premature_escalation"

        elif t == ActionType.CREATE_TICKET:
            r.compliance_delta = 0.05
            r.info = "ticket_created"

        r.dwell_penalty = round(-0.005 * self.MINUTES_PER_STEP[self.task_id], 4)
        r.total = round(
            r.classification_delta + r.chain_coverage_delta + r.containment_delta
            + r.fp_penalty + r.compliance_delta + r.dwell_penalty,
            4,
        )
        r.total = max(-1.0, min(1.0, r.total))
        return r

    # ------------------------------------------------------------------
    # Done condition
    # ------------------------------------------------------------------

    def _check_done(self, s: EnvState) -> bool:
        if s.step >= self.MAX_STEPS[self.task_id]:
            return True
        if s.attack_chain and s.attack_chain.exfiltration_complete:
            return True
        all_alerts = s.observation.active_alerts + s.observation.acknowledged_alerts
        closed_statuses = {AlertStatus.CONTAINED, AlertStatus.CLOSED, AlertStatus.FALSE_POSITIVE}
        if all_alerts and all(a.status in closed_statuses for a in all_alerts):
            return True
        return False

    # ------------------------------------------------------------------
    # Lookup helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _find_alert(alert_id: Optional[str], s: EnvState) -> Optional[SIEMAlert]:
        if not alert_id:
            return None
        for a in s.observation.active_alerts + s.observation.acknowledged_alerts:
            if a.alert_id == alert_id:
                return a
        return None

    @staticmethod
    def _find_host(host_id: Optional[str], s: EnvState) -> Optional[NetworkHost]:
        if not host_id:
            return None
        for h in s.observation.hosts:
            if h.host_id == host_id:
                return h
        return None