Spaces:
Sleeping
Sleeping
| """ | |
| Stratagem — Cybersecurity Incident Response Inference Script | |
| =================================== | |
| MANDATORY env vars: API_BASE_URL, MODEL_NAME, HF_TOKEN | |
| STDOUT FORMAT: [START], [STEP], [END] | |
| """ | |
| import asyncio | |
| import json | |
| import os | |
| import re | |
| import textwrap | |
| from typing import Any, Dict, List, Optional | |
| from openai import OpenAI | |
| from client import StratagemEnv | |
| from models import IncidentAction, ACTION_NAMES, NUM_ACTIONS, SYSTEM_NAMES | |
| # --------------------------------------------------------------------------- | |
| # Configuration | |
| # --------------------------------------------------------------------------- | |
| IMAGE_NAME = os.getenv("IMAGE_NAME") | |
| API_KEY = os.getenv("HF_TOKEN") or os.getenv("API_KEY") or os.getenv("OPENAI_API_KEY") | |
| API_BASE_URL = os.getenv("API_BASE_URL") or "https://router.huggingface.co/v1" | |
| MODEL_NAME = os.getenv("MODEL_NAME") or "Qwen/Qwen2.5-72B-Instruct" | |
| BENCHMARK = "stratagem" | |
| MAX_STEPS = 12 | |
| TEMPERATURE = 0.3 | |
| MAX_TOKENS = 400 | |
| TASKS = ["easy_1", "medium_1", "hard_1"] | |
| # --------------------------------------------------------------------------- | |
| # Structured stdout logging | |
| # --------------------------------------------------------------------------- | |
| def log_start(task: str, env: str, model: str) -> None: | |
| print(f"[START] task={task} env={env} model={model}", flush=True) | |
| def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None: | |
| error_val = error if error else "null" | |
| done_val = str(done).lower() | |
| print(f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}", flush=True) | |
| def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None: | |
| rewards_str = ",".join(f"{r:.2f}" for r in rewards) | |
| print(f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}", flush=True) | |
| # --------------------------------------------------------------------------- | |
| # System prompt | |
| # --------------------------------------------------------------------------- | |
| SYSTEM_PROMPT = textwrap.dedent("""\ | |
| You are an expert Incident Commander responding to a live cyberattack on a corporate network. | |
| ## Network (8 systems, indexed 0-7): | |
| 0: web_server 1: app_server 2: database 3: file_server | |
| 4: email_server 5: workstations 6: backup_server 7: firewall | |
| ## Actions (pick action 0-9 and target system 0-7): | |
| 0: investigate_system — Reveals true state of target. Costs stamina. Takes time. | |
| 1: isolate_system — Cuts target from network. Stops attacker BUT kills service. | |
| 2: patch_vulnerability — Fixes vuln on target. Slow, may clean compromised systems. | |
| 3: restore_from_backup — Restores target from backup. DANGER: backup may be compromised. | |
| 4: analyze_alerts — Deep analysis of alert queue. Reveals true/false positives. | |
| 5: deploy_monitoring — Adds sensors to target + neighbors. Improves future detection. | |
| 6: escalate_to_management — Gets resources but adds scrutiny pressure. | |
| 7: block_external_traffic — Stops ALL outbound connections. Kills exfiltration + services. | |
| 8: hunt_threat — Proactively search target for attacker indicators. | |
| 9: coordinate_team — Rest and regroup. Recovers stamina but wastes an hour. | |
| ## Key dynamics: | |
| - Attacker moves laterally through connected systems each hour | |
| - Attacker exfiltrates data from database, file_server, email_server, backup_server | |
| - Compromise is UNKNOWN until you investigate or hunt a system | |
| - Alerts may be FALSE POSITIVES — analyze_alerts reveals which are real | |
| - Team stamina depletes with actions; exhausted team is less effective | |
| - Restoring from a compromised backup re-infects the target! | |
| - Blocking external traffic stops exfiltration but disrupts all services | |
| ## Scoring: | |
| 35% data protection + 25% containment + 20% business continuity + 10% forensics + 10% team health | |
| ## Strategy tips: | |
| - INVESTIGATE before acting blindly | |
| - Prioritize isolating CONFIRMED compromised systems adjacent to critical data | |
| - Don't isolate everything — you need services running | |
| - Deploy monitoring EARLY for better future alerts | |
| - Watch your team stamina — coordinate_team recovers it | |
| Respond with ONLY: {"action": <0-9>, "target": <0-7>, "reasoning": "<brief>"} | |
| """) | |
| # --------------------------------------------------------------------------- | |
| # Observation formatting | |
| # --------------------------------------------------------------------------- | |
| def format_observation(obs: dict, step: int, history: List[str]) -> str: | |
| parts = [] | |
| if step == 0: | |
| desc = obs.get("task_description", "") | |
| if desc: | |
| parts.append(f"## Incident Brief\n{desc}\n") | |
| parts.append(f"## Hour {obs.get('hour', 0)} Status (Hours remaining: {obs.get('hours_remaining', 12)})") | |
| parts.append(f"- Breach severity: {obs.get('estimated_breach_severity', 'unknown')}") | |
| parts.append(f"- Data at risk: {obs.get('estimated_data_at_risk', 0):.0%}") | |
| parts.append(f"- Services disrupted: {obs.get('services_disrupted', 0)}/{obs.get('services_total', 4)}") | |
| parts.append(f"- Team stamina: {obs.get('team_stamina', 1.0):.0%}") | |
| parts.append(f"- External traffic blocked: {obs.get('external_blocked', False)}") | |
| parts.append(f"- Management escalated: {obs.get('management_escalated', False)}") | |
| # System statuses | |
| systems = obs.get("systems_visible", []) | |
| if systems: | |
| parts.append("\n## Systems") | |
| for s in systems: | |
| status_parts = [] | |
| comp = s.get("compromised", "unknown") | |
| if comp == "unknown": | |
| status_parts.append("compromise=?") | |
| else: | |
| status_parts.append(f"compromised={'YES' if comp else 'no'}") | |
| if s.get("isolated"): | |
| status_parts.append("ISOLATED") | |
| if s.get("investigated"): | |
| status_parts.append("investigated") | |
| if s.get("patched"): | |
| status_parts.append("patched") | |
| status_parts.append(f"integrity={s.get('integrity', 1.0):.0%}") | |
| status_parts.append(f"monitoring={s.get('monitoring_level', 0)}") | |
| parts.append(f" [{SYSTEM_NAMES.index(s['name'])}] {s['name']:16s} | {', '.join(status_parts)}") | |
| # Alerts | |
| alerts = obs.get("alert_queue", []) | |
| if alerts: | |
| parts.append("\n## Recent Alerts") | |
| for a in alerts[-4:]: | |
| confirmed = a.get("confirmed", "") | |
| conf_str = f" [{'CONFIRMED' if confirmed else 'FALSE POSITIVE'}]" if confirmed != "" else "" | |
| parts.append(f" [{a.get('severity', '?'):8s}] {a.get('message', '')}{conf_str}") | |
| if history: | |
| parts.append("\n## Your recent actions") | |
| for h in history[-3:]: | |
| parts.append(f" {h}") | |
| parts.append('\nRespond: {"action": <0-9>, "target": <0-7>, "reasoning": "..."}') | |
| return "\n".join(parts) | |
| def parse_response(text: str) -> tuple[int, int]: | |
| """Extract action and target from LLM response.""" | |
| json_match = re.search(r'\{[^}]*"action"\s*:\s*(\d)[^}]*\}', text) | |
| if json_match: | |
| try: | |
| data = json.loads(json_match.group()) | |
| action = int(data.get("action", 9)) | |
| target = int(data.get("target", 0)) | |
| if 0 <= action < NUM_ACTIONS and 0 <= target < len(SYSTEM_NAMES): | |
| return action, target | |
| except (json.JSONDecodeError, KeyError, ValueError): | |
| pass | |
| return 9, 0 # fallback: coordinate_team | |
| # --------------------------------------------------------------------------- | |
| # Run one task | |
| # --------------------------------------------------------------------------- | |
| async def run_task(env: StratagemEnv, task_id: str, client: OpenAI) -> float: | |
| history: List[str] = [] | |
| messages: List[Dict[str, str]] = [] | |
| rewards: List[float] = [] | |
| steps_taken = 0 | |
| score = 0.0 | |
| success = False | |
| log_start(task=task_id, env=BENCHMARK, model=MODEL_NAME) | |
| try: | |
| result = await env.reset(task_id=task_id) | |
| obs = result.observation.model_dump() | |
| for step in range(1, MAX_STEPS + 1): | |
| if result.done: | |
| break | |
| user_msg = format_observation(obs, step - 1, history) | |
| messages.append({"role": "user", "content": user_msg}) | |
| try: | |
| completion = client.chat.completions.create( | |
| model=MODEL_NAME, | |
| messages=[{"role": "system", "content": SYSTEM_PROMPT}] + messages, | |
| max_tokens=MAX_TOKENS, | |
| temperature=TEMPERATURE, | |
| stream=False, | |
| ) | |
| llm_text = (completion.choices[0].message.content or "").strip() | |
| except Exception as exc: | |
| print(f"[DEBUG] LLM error: {exc}", flush=True) | |
| llm_text = '{"action": 9, "target": 0, "reasoning": "API error fallback"}' | |
| messages.append({"role": "assistant", "content": llm_text}) | |
| action_idx, target_idx = parse_response(llm_text) | |
| action_name = f"{ACTION_NAMES.get(action_idx, str(action_idx))}({SYSTEM_NAMES[target_idx]})" | |
| result = await env.step(IncidentAction(action=action_idx, target_system=target_idx)) | |
| obs = result.observation.model_dump() | |
| reward = result.reward or 0.0 | |
| done = result.done | |
| error = obs.get("metadata", {}).get("error") if isinstance(obs.get("metadata"), dict) else None | |
| rewards.append(reward) | |
| steps_taken = step | |
| log_step(step=step, action=action_name, reward=reward, done=done, error=error) | |
| history.append(f"Hour {step}: {action_name} -> reward {reward:+.2f}") | |
| if done: | |
| meta = result.observation.metadata or {} | |
| score = meta.get("comparison_score", 0.5) | |
| score = min(max(score, 0.0), 1.0) | |
| success = score >= 0.5 | |
| break | |
| if not result.done: | |
| score = 0.5 | |
| success = True | |
| except Exception as exc: | |
| print(f"[DEBUG] Task {task_id} error: {exc}", flush=True) | |
| score = 0.0 | |
| success = False | |
| finally: | |
| log_end(success=success, steps=steps_taken, score=score, rewards=rewards) | |
| return score | |
| # --------------------------------------------------------------------------- | |
| # Main | |
| # --------------------------------------------------------------------------- | |
| async def main() -> None: | |
| client = OpenAI(base_url=API_BASE_URL, api_key=API_KEY) | |
| env = await StratagemEnv.from_docker_image(IMAGE_NAME) | |
| try: | |
| for task_id in TASKS: | |
| await run_task(env, task_id, client) | |
| finally: | |
| try: | |
| await env.close() | |
| except Exception as e: | |
| print(f"[DEBUG] env.close() error: {e}", flush=True) | |
| if __name__ == "__main__": | |
| asyncio.run(main()) | |