"""Exploit chain models for the redteam module."""
from __future__ import annotations

from datetime import UTC, datetime
from enum import Enum
from typing import Literal

from pydantic import BaseModel, Field


class GoalType(str, Enum):
    PROMPT_DRIVEN_THREAT = "PROMPT_DRIVEN_THREAT"
    POLICY_VIOLATION = "POLICY_VIOLATION"
    DATA_EXFILTRATION = "DATA_EXFILTRATION"
    PRIVILEGE_ESCALATION = "PRIVILEGE_ESCALATION"
    TOOL_ABUSE = "TOOL_ABUSE"
    MCP_TOXIC_FLOW = "MCP_TOXIC_FLOW"
    API_ATTACK = "API_ATTACK"


class ScenarioType(str, Enum):
    SYSTEM_PROMPT_EXTRACTION = "SYSTEM_PROMPT_EXTRACTION"
    GUARDRAIL_BYPASS = "GUARDRAIL_BYPASS"
    INDIRECT_INJECTION = "INDIRECT_INJECTION"
    MULTI_TURN_REDIRECTION = "MULTI_TURN_REDIRECTION"
    RESTRICTED_ACTION = "RESTRICTED_ACTION"
    HITL_BYPASS = "HITL_BYPASS"
    DIRECT_PII_EXTRACTION = "DIRECT_PII_EXTRACTION"
    CROSS_TENANT_EXFILTRATION = "CROSS_TENANT_EXFILTRATION"
    COVERT_ENCODING = "COVERT_ENCODING"
    RAG_POISONING = "RAG_POISONING"
    PRIVILEGE_CHAIN = "PRIVILEGE_CHAIN"
    SQL_INJECTION = "SQL_INJECTION"
    SSRF = "SSRF"
    MASS_ASSIGNMENT = "MASS_ASSIGNMENT"
    MCP_WEB_FETCH = "MCP_WEB_FETCH"
    MCP_EMAIL = "MCP_EMAIL"
    MCP_GITHUB = "MCP_GITHUB"
    MCP_RAG = "MCP_RAG"
    AUTH_BYPASS = "AUTH_BYPASS"
    IDOR = "IDOR"
    CONTEXT_FLOODING = "CONTEXT_FLOODING"
    STRUCTURAL_INJECTION = "STRUCTURAL_INJECTION"
    MCP_TOOL_INJECTION = "MCP_TOOL_INJECTION"


# Sentinel success_signal value for HTTP-direct steps: success = any 2xx status code.
HTTP_2XX_SENTINEL = "HTTP_2XX"


class ExploitStep(BaseModel):
    step_id: str
    step_type: Literal["INJECT", "PIVOT", "INVOKE", "SCAN", "EVALUATE", "OBSERVE"]
    description: str
    payload: str
    target_node_id: str = ""
    depends_on: list[str] = Field(default_factory=list)
    success_signal: str = ""
    on_failure: Literal["abort", "mutate", "skip"] = "skip"
    result: str | None = None
    succeeded: bool | None = None
    # When True, the executor will use LLMResponseEvaluator to determine step success
    # instead of (or in addition to) the keyword success_signal.  Set for complex
    # scenarios (e.g. HITL_BYPASS) where keyword matching is unreliable.
    use_llm_eval: bool = False
    # When True, a keyword success_signal only counts if the HTTP response status is
    # 2xx.  Prevents false positives when frameworks (e.g. FastAPI) echo the request
    # body verbatim inside 4xx validation error responses.
    success_requires_2xx: bool = False
    # HTTP-direct attack fields — when set, the step bypasses the chat endpoint and
    # sends a raw HTTP request to target_path instead.
    target_path: str | None = None
    http_method: str = "POST"
    http_params: dict[str, str] = Field(default_factory=dict)
    http_body: dict | None = None


class ExploitChain(BaseModel):
    chain_id: str
    goal_type: GoalType
    scenario_type: ScenarioType
    steps: list[ExploitStep] = Field(default_factory=list)
    pre_score: float = 0.0
    sbom_path: list[str] = Field(default_factory=list)
    policy_clauses: list[str] = Field(default_factory=list)
    mitre_atlas_technique: str | None = None
    owasp_asi_ref: str | None = None
    owasp_llm_ref: str | None = None
    status: Literal["pending", "running", "completed", "aborted", "failed"] = "pending"
    created_at: datetime = Field(default_factory=lambda: datetime.now(UTC))
