server/app.py

"""FastAPI server for the SOC triage environment."""

from __future__ import annotations

import inspect
import os
from typing import Any

from fastapi import Body, HTTPException
from pydantic import BaseModel, Field

try:
    from openenv.core.env_server.http_server import create_app
except Exception:
    from openenv.core.env_server import create_app

try:
    from ..models import TriageAction, TriageObservation
    from .graders import grade_easy, grade_hard, grade_medium
    from .soc_triage_env import SOCTriageEnv
    from .tasks import TASKS
except ImportError:
    from soc_triage_env.models import TriageAction, TriageObservation
    from soc_triage_env.server.graders import grade_easy, grade_hard, grade_medium
    from soc_triage_env.server.soc_triage_env import SOCTriageEnv
    from soc_triage_env.server.tasks import TASKS


class GraderRequest(BaseModel):
    task_id: str
    action: TriageAction
    ground_truth: dict[str, Any] | None = None


class BaselineRequest(BaseModel):
    provider: str = "blaxel"
    model: str = "sandbox-openai"
    fallback_provider: str = "cerebras"
    fallback_model: str = "llama3.1-8b"
    episodes_per_task: int = Field(default=1, ge=1, le=5)


import json
import time
from pathlib import Path

def _build_app() -> Any:
    sig = inspect.signature(create_app)
    kwargs: dict[str, Any] = {
        "env_name": "soc_triage_env",
        "max_concurrent_envs": 10,
    }
    if "max_concurrent_envs" not in sig.parameters:
        kwargs.pop("max_concurrent_envs", None)
    return create_app(SOCTriageEnv, TriageAction, TriageObservation, **kwargs)


app = _build_app()
_grader_env = SOCTriageEnv()

LOG_FILE = Path(__file__).parent / "validator_tests.log"

@app.middleware("http")
async def log_requests(request, call_next):
    start_time = time.time()
    body_bytes = await request.body()
    
    async def receive():
        return {"type": "http.request", "body": body_bytes}
    request._receive = receive
    
    response = await call_next(request)
    process_time = time.time() - start_time
    
    body_str = body_bytes.decode('utf-8', errors='ignore')
    log_entry = {
        "timestamp": time.time(),
        "method": request.method,
        "url": str(request.url),
        "status": response.status_code,
        "latency_sec": round(process_time, 4),
        "request_body": body_str
    }
    
    with open(LOG_FILE, "a") as f:
        f.write(json.dumps(log_entry) + "\n")
        
    return response


@app.get("/")
def root() -> dict[str, Any]:
    return {
        "name": "SOC Triage OpenEnv",
        "status": "ok",
        "mode": "interactive",
        "endpoints": [
            "/health",
            "/reset",
            "/step",
            "/state",
            "/schema",
            "/tasks",
            "/grader",
            "/baseline",
            "/ws",
            "/web",
            "/logs",
        ],
    }

@app.get("/logs")
def get_logs() -> dict[str, Any]:
    if not LOG_FILE.exists():
        return {"logs": []}
    logs = []
    with open(LOG_FILE, "r") as f:
        lines = f.readlines()
        for line in lines[-100:]:
            try:
                logs.append(json.loads(line))
            except Exception:
                pass
    return {"logs": logs}

@app.get("/tasks")
def tasks() -> dict[str, Any]:
    return {"tasks": TASKS}


@app.post("/grader")
def grader(payload: GraderRequest) -> dict[str, Any]:
    task_id = payload.task_id.strip().lower()
    if task_id not in TASKS:
        raise HTTPException(status_code=400, detail=f"Unsupported task_id: {task_id}")

    ground_truth = payload.ground_truth or _grader_env.ground_truth
    classification = _classification_from_action(payload.action)

    if task_id == "easy":
        score = grade_easy(classification, str(ground_truth.get("severity", "benign")))
    elif task_id == "medium":
        score = grade_medium(_parse_ids(classification), list(ground_truth.get("ranking", [])))
    else:
        score = grade_hard(_parse_ids(classification), list(ground_truth.get("kill_chain", [])))

    return {"task_id": task_id, "score": round(score, 4)}


@app.post("/baseline")
def baseline(payload: BaselineRequest = Body(default_factory=BaselineRequest)) -> dict[str, Any]:
    """Run provider baseline if keys are set, else return heuristic fallback."""
    provider = payload.provider.lower().strip()
    if provider not in {"openai", "cerebras", "blaxel"}:
        raise HTTPException(status_code=400, detail="provider must be openai, cerebras, or blaxel")

    try:
        from baseline import run_baseline_with_fallback_sync

        mode, scores, warning = run_baseline_with_fallback_sync(
            provider=provider,
            model=payload.model,
            episodes_per_task=payload.episodes_per_task,
            fallback_provider=payload.fallback_provider,
            fallback_model=payload.fallback_model,
        )
        response = {"mode": mode, "model": payload.model, "scores": scores}
        if warning:
            response["warning"] = warning
        return response
    except Exception as exc:
        return {
            "mode": "heuristic",
            "warning": f"Provider baseline unavailable (primary and fallback failed: {exc}).",
            "scores": _heuristic_baseline(),
        }


def _heuristic_baseline() -> dict[str, float]:
    local_env = SOCTriageEnv()
    scores: dict[str, float] = {}

    verdicts = {
        "easy": TriageAction(
            tool_name="submit_verdict",
            classification="high",
            recommended_action="investigate",
            reasoning="Heuristic baseline for easy task.",
        ),
        "medium": TriageAction(
            tool_name="submit_verdict",
            classification="MED-C,MED-E,MED-D,MED-A,MED-B",
            recommended_action="investigate",
            reasoning="Heuristic ranking using known priority signal order.",
        ),
        "hard": TriageAction(
            tool_name="submit_verdict",
            classification="H-01,H-03,H-05,H-07,H-11",
            recommended_action="contain",
            reasoning="Heuristic kill-chain pattern match.",
        ),
    }

    for task_id, verdict in verdicts.items():
        local_env.reset(task_id=task_id)
        # Perform one investigative action before submission to use multi-turn mechanics.
        local_env.step(TriageAction(tool_name="query_siem", tool_args={"query": "suspicious"}))
        obs = local_env.step(verdict)
        scores[task_id] = round(obs.reward, 4)
    return scores


def _classification_from_action(action: TriageAction) -> str:
    if action.classification:
        return action.classification
    if isinstance(action.tool_args, dict):
        value = action.tool_args.get("classification", "")
        return str(value).strip()
    return ""


def _parse_ids(value: str) -> list[str]:
    return [x.strip() for x in value.split(",") if x.strip()]


def main(host: str = "0.0.0.0", port: int | None = None) -> None:
    """Run the API with uvicorn for local and validator execution."""
    import uvicorn

    resolved_port = port if port is not None else int(os.getenv("API_PORT", "8000"))
    uvicorn.run(app, host=host, port=resolved_port)


if __name__ == "__main__":
    main()