baseline.py

"""
Baseline inference script — uses OpenAI API client to run an LLM agent
against all 3 tasks and produces reproducible scores.

Usage
-----
    export OPENAI_API_KEY="sk-..."
    export ENV_URL="http://localhost:8000"   # optional, defaults to localhost
    python baseline.py

The script prints per-task and aggregate scores, then writes results to
baseline_results.json for reproducibility.
"""

from __future__ import annotations

import asyncio
import json
import os
import re
import sys
from datetime import datetime
from typing import Dict, List, Optional

from openai import AsyncOpenAI

from client import IncidentResponseEnv, StepResult
from graders import GradeResult, grade
from models import IncidentAction, IncidentObservation

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

API_KEY = os.getenv("API_KEY", "")
ENV_URL        = os.getenv("API_BASE_URL", "http://localhost:8000")
MODEL_NAME     = os.getenv("BASELINE_MODEL", "gpt-4o-mini")
NUM_EPISODES   = int(os.getenv("NUM_EPISODES", "3"))

if not API_KEY:
    print("ERROR: API_KEY environment variable is not set.", file=sys.stderr)
    sys.exit(1)

client = AsyncOpenAI(api_key= API_KEY)

# ---------------------------------------------------------------------------
# System prompt
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = """\
You are an expert cybersecurity incident responder AI agent.
You will be given the current state of a simulated company under cyber attack.
Your goal is to investigate logs and alerts, identify the threat, and mitigate it before the system is compromised.

## Available Actions (one per step):
- analyze_log(log_id)      — Examine a specific log entry for more detail
- trace_user(user_id)      — Investigate a user's activity history
- block_ip(ip_address)     — Block a suspicious IP address
- disable_account(user_id) — Disable a compromised user account
- restart_service(service) — Restart a compromised or degraded service
- ignore                   — Take no action (penalised — attack progresses!)

## Response Format (STRICT — machine-parsed):
You MUST respond ONLY with a JSON object like:
{"action_type": "analyze_log", "target": "L001"}

Valid action_type values: analyze_log, trace_user, block_ip, disable_account, restart_service, ignore

## Strategy:
1. First investigate (analyze_log, trace_user) to gather evidence.
2. Then act decisively on confirmed threats (block_ip, disable_account, restart_service).
3. Never block/disable unless you have strong evidence — wrong actions cost points.
4. Speed matters — faster resolution earns a bonus.
"""

# ---------------------------------------------------------------------------
# LLM-driven agent
# ---------------------------------------------------------------------------

def _format_observation(obs: IncidentObservation) -> str:
    logs_txt    = "\n".join(f"  [{l.log_id}] {l.timestamp} — {l.message}" for l in obs.logs)
    alerts_txt  = "\n".join(f"  [{a.severity.upper()}] {a.message}" for a in obs.alerts)
    services_txt = "\n".join(f"  {s.name}: {s.status}" for s in obs.services)
    users_txt   = ", ".join(obs.active_users)

    return f"""\
=== Incident Response Dashboard (Step {obs.step_count}) ===

LOGS:
{logs_txt}

ALERTS:
{alerts_txt}

SERVICES:
{services_txt}

ACTIVE USERS: {users_txt}

Last action result: {obs.last_action_result}
"""


def _parse_llm_response(text: str) -> IncidentAction:
    """Extract JSON action from LLM output. Falls back to ignore on parse failure."""
    # Try direct JSON parse
    stripped = text.strip()
    try:
        data = json.loads(stripped)
        return IncidentAction(**data)
    except Exception:
        pass

    # Try extracting JSON from markdown code block
    match = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", stripped, re.DOTALL)
    if match:
        try:
            data = json.loads(match.group(1))
            return IncidentAction(**data)
        except Exception:
            pass

    # Try finding raw JSON object in text
    match = re.search(r"\{[^{}]+\}", stripped)
    if match:
        try:
            data = json.loads(match.group(0))
            return IncidentAction(**data)
        except Exception:
            pass

    # Fallback
    print(f"  [WARN] Could not parse LLM output: {text[:120]!r} — defaulting to ignore")
    return IncidentAction(action_type="ignore", target="")


async def _llm_agent_fn(
    obs: IncidentObservation,
    history: List[StepResult],
    conversation: List[Dict],
) -> IncidentAction:
    """Call OpenAI API and return the next action."""
    user_msg = _format_observation(obs)

    conversation.append({"role": "user", "content": user_msg})

    response = await client.chat.completions.create(
        model=MODEL_NAME,
        messages=[{"role": "system", "content": SYSTEM_PROMPT}] + conversation,
        max_tokens=256,
        temperature=0.0,  # deterministic for reproducibility
    )

    assistant_text = response.choices[0].message.content or ""
    conversation.append({"role": "assistant", "content": assistant_text})

    return _parse_llm_response(assistant_text)


# ---------------------------------------------------------------------------
# Episode runner
# ---------------------------------------------------------------------------

async def _run_llm_episode(task_id: int) -> GradeResult:
    conversation: List[Dict] = []
    last_result: Optional[StepResult] = None

    async with IncidentResponseEnv(base_url=ENV_URL) as env:
        obs = await env.reset(task_id=task_id)
        history: List[StepResult] = []

        for step in range(25):  # safety cap
            action = await _llm_agent_fn(obs, history, conversation)
            print(f"    step {step + 1:02d}: {action.action_type}({action.target!r})", end="")

            result = await env.step(action)
            print(f"  → reward={result.reward:+.2f}")

            history.append(result)
            last_result = result
            obs = result.observation

            if result.done:
                break

    info = last_result.info if last_result else {}
    return GradeResult(
        task_id=task_id,
        score=float(info.get("final_score", 0.0)),
        success=bool(info.get("success", False)),
        compromise=bool(info.get("compromise", False)),
        steps_taken=obs.step_count,
        wrong_actions=int(info.get("wrong_actions", 0)),
        notes=obs.last_action_result,
    )


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

async def main():
    print("=" * 65)
    print(f"Incident Response Env — Baseline ({MODEL_NAME})")
    print(f"Environment: {ENV_URL}")
    print(f"Episodes per task: {NUM_EPISODES}")
    print("=" * 65)

    all_scores: List[float] = []
    output: Dict = {
        "model":        MODEL_NAME,
        "env_url":      ENV_URL,
        "timestamp":    datetime.utcnow().isoformat(),
        "tasks":        {},
    }

    for task_id in [1, 2, 3]:
        task_names = {
            1: "Brute Force (Easy)",
            2: "Suspicious Login (Medium)",
            3: "Multi-Stage APT (Hard)",
        }
        print(f"\n--- Task {task_id}: {task_names[task_id]} ---")

        episode_results = []
        for ep in range(NUM_EPISODES):
            print(f"  Episode {ep + 1}/{NUM_EPISODES}:")
            result = await _run_llm_episode(task_id)
            episode_results.append(result)
            print(f"  → {result}")

        scores = [r.score for r in episode_results]
        mean   = sum(scores) / len(scores)
        all_scores.extend(scores)

        output["tasks"][str(task_id)] = {
            "mean_score":   round(mean, 4),
            "success_rate": round(sum(1 for r in episode_results if r.success) / NUM_EPISODES, 4),
            "episodes":     [
                {"score": r.score, "success": r.success, "steps": r.steps_taken}
                for r in episode_results
            ],
        }
        print(f"  Task {task_id} mean score: {mean:.4f}")

    overall = sum(all_scores) / len(all_scores)
    output["overall_mean_score"] = round(overall, 4)

    print(f"\n{'=' * 65}")
    print(f"Overall mean score: {overall:.4f}")
    print("=" * 65)

    out_path = "baseline_results.json"
    with open(out_path, "w") as f:
        json.dump(output, f, indent=2)
    print(f"\nResults saved to {out_path}")


if __name__ == "__main__":
    asyncio.run(main())