models.py

"""
Data models for the ForensicShell OpenEnv environment.

ForensicShell simulates a pre-seeded "breached" Linux box. The agent investigates
via structured read-only actions (list_dir, read_file, grep, stat) and ultimately
submits a ForensicReport. A deterministic grader scores the report against hidden
ground truth per task (easy/medium/hard).
"""

from typing import List, Literal, Optional

from openenv.core.env_server.types import Action, Observation
from pydantic import BaseModel, Field


class TimelineEvent(BaseModel):
    """One step in the attacker's kill chain (used in hard task)."""

    phase: Literal["login", "recon", "privesc", "persistence", "exfil"]
    detail: str = Field(default="", description="Short description of the event")


class ForensicReport(BaseModel):
    """The agent's final investigation report. Submitted via action_type='submit_report'."""

    compromised_user: Optional[str] = Field(
        default=None, description="Username of the compromised account"
    )
    initial_ip: Optional[str] = Field(
        default=None, description="Source IP of the initial login"
    )
    modified_files: List[str] = Field(
        default_factory=list,
        description="Absolute paths of files modified after compromise",
    )
    backdoor_sha256: Optional[str] = Field(
        default=None,
        description="SHA256 of the attacker-dropped backdoor binary (lowercase hex)",
    )
    timeline: List[TimelineEvent] = Field(
        default_factory=list,
        description="Ordered attacker kill chain (login→recon→privesc→persistence→exfil)",
    )


class ForensicShellAction(Action):
    """Agent action. Use action_type to pick the verb; set only the fields that verb needs."""

    action_type: Literal[
        "list_dir", "read_file", "grep", "stat", "find", "submit_report"
    ] = Field(..., description="Which verb to execute")
    path: Optional[str] = Field(
        default=None, description="Target path for list_dir / read_file / grep / stat"
    )
    pattern: Optional[str] = Field(
        default=None, description="Substring pattern for grep"
    )
    max_bytes: int = Field(
        default=2048, description="Max bytes to return from read_file (truncated)"
    )
    report: Optional[ForensicReport] = Field(
        default=None, description="Investigation report (only for action_type='submit_report')"
    )


class ForensicShellObservation(Observation):
    """Result of the agent's last action."""

    output: str = Field(default="", description="Human-readable result of the action")
    task_id: str = Field(default="", description="Current task identifier")
    task_description: str = Field(
        default="", description="What the agent is being asked to investigate"
    )
    steps_remaining: int = Field(
        default=0, description="How many more actions allowed this episode"
    )
    action_error: Optional[str] = Field(
        default=None, description="Error message if the action failed; None otherwise"
    )