arushsingh/soc-analyst-env

1"""2Data models for the SOC Analyst Environment.3 4Defines the Action, Observation, and internal data types for a Security5Operations Center analyst triage and investigation simulation.6"""7 8from enum import Enum9from typing import Any, Dict, List, Optional10 11from openenv.core.env_server.types import Action, Observation, State12from pydantic import Field13 14 15# --- Enums ---16 17 18class TaskType(str, Enum):19    PHISHING_TRIAGE = "phishing_triage"20    MALWARE_INVESTIGATION = "malware_investigation"21    APT_DETECTION = "apt_detection"22 23 24class Severity(str, Enum):25    LOW = "low"26    MEDIUM = "medium"27    HIGH = "high"28    CRITICAL = "critical"29 30 31class AlertVerdict(str, Enum):32    TRUE_POSITIVE = "true_positive"33    FALSE_POSITIVE = "false_positive"34    BENIGN = "benign"35    SUSPICIOUS = "suspicious"36 37 38class RemediationAction(str, Enum):39    BLOCK_IP = "block_ip"40    BLOCK_DOMAIN = "block_domain"41    QUARANTINE_HOST = "quarantine_host"42    QUARANTINE_EMAIL = "quarantine_email"43    DISABLE_ACCOUNT = "disable_account"44    ESCALATE_TO_TIER2 = "escalate_to_tier2"45    NO_ACTION = "no_action"46 47 48class LogSource(str, Enum):49    FIREWALL = "firewall"50    IDS_IPS = "ids_ips"51    ENDPOINT = "endpoint"52    EMAIL_GATEWAY = "email_gateway"53    AUTH_LOGS = "auth_logs"54    DNS_LOGS = "dns_logs"55    PROXY_LOGS = "proxy_logs"56 57 58class KillChainPhase(str, Enum):59    RECONNAISSANCE = "reconnaissance"60    INITIAL_ACCESS = "initial_access"61    EXECUTION = "execution"62    PERSISTENCE = "persistence"63    PRIVILEGE_ESCALATION = "privilege_escalation"64    LATERAL_MOVEMENT = "lateral_movement"65    COLLECTION = "collection"66    EXFILTRATION = "exfiltration"67    COMMAND_AND_CONTROL = "command_and_control"68 69 70# --- Action / Observation ---71 72VALID_ACTIONS = [73    "get_alert_queue",74    "examine_alert",75    "examine_email",76    "query_logs",77    "check_threat_intel",78    "check_url_reputation",79    "check_file_hash",80    "correlate_events",81    "get_endpoint_details",82    "classify_alert",83    "take_remediation",84    "submit_incident_report",85]86 87 88class SOCAction(Action):89    """Action for the SOC Analyst environment.90 91    The agent selects an action_type and provides parameters in params.92    This design maps naturally to LLM function-calling.93    """94 95    action_type: str = Field(96        ...,97        description=(98            "The investigation action to take. One of: "99            + ", ".join(VALID_ACTIONS)100        ),101    )102    params: Dict[str, Any] = Field(103        default_factory=dict,104        description="Parameters for the selected action",105    )106 107 108class SOCObservation(Observation):109    """Observation returned by the SOC Analyst environment.110 111    Contains a human-readable message and structured data from the112    last action, plus episode context.113    """114 115    message: str = Field(default="", description="Human-readable result description")116    data: Dict[str, Any] = Field(117        default_factory=dict, description="Structured result data"118    )119    task_type: str = Field(default="", description="Current task type")120    step_number: int = Field(default=0, description="Current step number")121    max_steps: int = Field(default=0, description="Maximum steps for this episode")122    available_actions: List[str] = Field(123        default_factory=list, description="Actions the agent can take"124    )125 126 127# --- Extended State ---128 129 130class SOCState(State):131    """Tracks investigation progress across the episode."""132 133    task_type: str = ""134    scenario_id: str = ""135    alerts_examined: List[str] = Field(default_factory=list)136    logs_queried: List[str] = Field(default_factory=list)137    threat_intel_checked: List[str] = Field(default_factory=list)138    urls_checked: List[str] = Field(default_factory=list)139    hashes_checked: List[str] = Field(default_factory=list)140    endpoints_checked: List[str] = Field(default_factory=list)141    emails_examined: List[str] = Field(default_factory=list)142    correlations_run: int = 0143    verdicts_submitted: Dict[str, str] = Field(default_factory=dict)144    remediations_taken: List[Dict[str, str]] = Field(default_factory=list)145    incident_report: Optional[Dict[str, Any]] = None146    evidence_collected: List[str] = Field(default_factory=list)147 148 149# --- Internal Scenario Data Types ---150 151 152class LogEntry:153    """A single log entry from a security data source."""154 155    __slots__ = (156        "timestamp", "source", "source_ip", "dest_ip",157        "message", "raw",158    )159 160    def __init__(161        self,162        timestamp: str,163        source: str,164        message: str,165        source_ip: str = "",166        dest_ip: str = "",167        raw: Optional[Dict[str, Any]] = None,168    ) -> None:169        self.timestamp = timestamp170        self.source = source171        self.source_ip = source_ip172        self.dest_ip = dest_ip173        self.message = message174        self.raw = raw or {}175 176    def to_dict(self) -> Dict[str, Any]:177        result: Dict[str, Any] = {178            "timestamp": self.timestamp,179            "source": self.source,180            "message": self.message,181        }182        if self.source_ip:183            result["source_ip"] = self.source_ip184        if self.dest_ip:185            result["dest_ip"] = self.dest_ip186        if self.raw:187            result["details"] = self.raw188        return result189 190 191class EmailData:192    """An email record for phishing investigation."""193 194    __slots__ = (195        "email_id", "from_address", "to_address", "subject",196        "body", "headers", "urls", "attachments",197    )198 199    def __init__(200        self,201        email_id: str,202        from_address: str,203        to_address: str,204        subject: str,205        body: str,206        headers: Optional[Dict[str, str]] = None,207        urls: Optional[List[str]] = None,208        attachments: Optional[List[Dict[str, str]]] = None,209    ) -> None:210        self.email_id = email_id211        self.from_address = from_address212        self.to_address = to_address213        self.subject = subject214        self.body = body215        self.headers = headers or {}216        self.urls = urls or []217        self.attachments = attachments or []218 219    def to_dict(self) -> Dict[str, Any]:220        return {221            "email_id": self.email_id,222            "from": self.from_address,223            "to": self.to_address,224            "subject": self.subject,225            "body": self.body,226            "headers": self.headers,227            "urls": self.urls,228            "attachments": self.attachments,229        }230 231 232class Alert:233    """A security alert in the SOC queue."""234 235    __slots__ = (236        "alert_id", "timestamp", "source", "title",237        "severity", "description", "indicators",238    )239 240    def __init__(241        self,242        alert_id: str,243        timestamp: str,244        source: str,245        title: str,246        severity: str,247        description: str,248        indicators: Optional[Dict[str, Any]] = None,249    ) -> None:250        self.alert_id = alert_id251        self.timestamp = timestamp252        self.source = source253        self.title = title254        self.severity = severity255        self.description = description256        self.indicators = indicators or {}257 258    def summary(self) -> Dict[str, str]:259        return {260            "alert_id": self.alert_id,261            "timestamp": self.timestamp,262            "severity": self.severity,263            "title": self.title,264            "source": self.source,265        }266 267    def to_dict(self) -> Dict[str, Any]:268        return {269            "alert_id": self.alert_id,270            "timestamp": self.timestamp,271            "source": self.source,272            "title": self.title,273            "severity": self.severity,274            "description": self.description,275            "indicators": self.indicators,276        }277 278 279class ThreatIntelEntry:280    """A threat intelligence database record."""281 282    __slots__ = (283        "indicator", "indicator_type", "threat_type",284        "confidence", "description", "tags",285    )286 287    def __init__(288        self,289        indicator: str,290        indicator_type: str,291        threat_type: str,292        confidence: float,293        description: str,294        tags: Optional[List[str]] = None,295    ) -> None:296        self.indicator = indicator297        self.indicator_type = indicator_type298        self.threat_type = threat_type299        self.confidence = confidence300        self.description = description301        self.tags = tags or []302 303    def to_dict(self) -> Dict[str, Any]:304        return {305            "indicator": self.indicator,306            "indicator_type": self.indicator_type,307            "threat_type": self.threat_type,308            "confidence": self.confidence,309            "description": self.description,310            "tags": self.tags,311        }312 313 314class Scenario:315    """Complete scenario definition for a task."""316 317    __slots__ = (318        "scenario_id", "task_type", "alerts", "log_database",319        "emails", "threat_intel_database", "url_reputation",320        "file_hashes", "endpoint_data", "ground_truth",321    )322 323    def __init__(324        self,325        scenario_id: str,326        task_type: str,327        alerts: List[Alert],328        log_database: Dict[str, List[LogEntry]],329        ground_truth: Dict[str, Any],330        emails: Optional[List[EmailData]] = None,331        threat_intel_database: Optional[List[ThreatIntelEntry]] = None,332        url_reputation: Optional[Dict[str, Dict[str, Any]]] = None,333        file_hashes: Optional[Dict[str, Dict[str, Any]]] = None,334        endpoint_data: Optional[Dict[str, Dict[str, Any]]] = None,335    ) -> None:336        self.scenario_id = scenario_id337        self.task_type = task_type338        self.alerts = alerts339        self.log_database = log_database340        self.emails = emails or []341        self.threat_intel_database = threat_intel_database or []342        self.url_reputation = url_reputation or {}343        self.file_hashes = file_hashes or {}344        self.endpoint_data = endpoint_data or {}345        self.ground_truth = ground_truth346