models.py

from __future__ import annotations

from typing import Any, Dict, List, Literal, Optional

from pydantic import Field

try:
    from core.env_server.types import Action, Observation, State
except ImportError:
    try:
        from openenv.core.env_server.types import Action, Observation, State
    except ImportError:
        from openenv_core.env_server.types import Action, Observation, State

ActionType = Literal["inspect_file", "submit_finding", "submit_final_report"]
VulnerabilityType = Literal[
    "sql_injection",
    "command_injection",
    "path_traversal",
    "weak_authentication",
    "insecure_deserialization",
    "ssrf",
    "hardcoded_secret",
    "xss",
]
Severity = Literal["low", "medium", "high", "critical"]


class CodeSecurityAction(Action):
    """Action sent by the agent during a security audit episode."""

    action_type: ActionType
    filename: Optional[str] = None
    line_start: Optional[int] = Field(default=None, ge=1)
    line_end: Optional[int] = Field(default=None, ge=1)
    vuln_type: Optional[VulnerabilityType] = None
    severity: Optional[Severity] = None
    confidence: float = Field(default=0.5, ge=0.0, le=1.0)
    evidence: str = ""
    summary: str = ""


class FindingRecord(State):
    """Stored record of one submitted finding."""

    finding_id: str
    filename: str
    line_start: int
    line_end: int
    vuln_type: str
    severity: str
    confidence: float
    evidence: str
    summary: str
    matched_vulnerability_id: Optional[str] = None
    component_score: float = 0.0


class CodeSecurityObservation(Observation):
    """Observation returned after reset() and step()."""

    task_id: str
    task_title: str
    difficulty: str
    objective: str
    instructions: str
    available_files: List[str] = Field(default_factory=list)
    focused_file: Optional[str] = None
    file_excerpt: str = ""
    findings_so_far: List[Dict[str, Any]] = Field(default_factory=list)
    steps_remaining: int = 0
    last_feedback: str = ""
    score_hint: float = Field(default=0.0, ge=0.0, le=1.0)


class CodeSecurityState(State):
    """Internal environment state for the current security auditing episode."""

    task_id: str = ""
    task_title: str = ""
    difficulty: str = ""
    objective: str = ""
    max_steps: int = 0
    inspected_files: List[str] = Field(default_factory=list)
    findings_submitted: List[FindingRecord] = Field(default_factory=list)
    matched_vulnerability_ids: List[str] = Field(default_factory=list)
    false_positive_count: int = 0
    duplicate_submission_count: int = 0
    quality_multiplier: float = 1.0
    final_score: Optional[float] = None