Update inference.py

inference.py

@@ -56,25 +56,40 @@ def warroom_step(att_type, att_res, scn_type, scn_tip, scn_res, rem_type, rem_ti
 
 # ── LLM: Single-agent action ──
 def get_single_action(task_id: str, obs: dict) -> dict:
-    prompt = f"""You are a cybersecurity AI agent.
+    terminal  = obs.get("terminal_output", "")
+    inventory = json.dumps(obs.get("inventory", []), indent=2)
+
+    rules = {
+        "easy-lockdown":       'command="lockdown", target_id="s3-vault"',
+        "easy-secrets":        'Step 1: command="audit", target_id="api-key-01" → Step 2: command="revoke", target_id="api-key-01" → Step 3: command="rotate", target_id="api-key-01"',
+        "medium-access":       'command="revoke_admin", target_id="user-dev-01"',
+        "medium-mfa":          'Step 1: command="audit", target_id="iam-users" → Then: command="enforce_mfa", target_id="admin-alice" (repeat for admin-bob and admin-carol)',
+        "hard-breach":         'Step 1: command="block_ip", target_id="attacker-ip" → Step 2: command="close_port", target_id="web-server"',
+        "critical-ransomware": 'Step 1: command="isolate", target_id="db-server" → Step 2: command="revoke_sessions", target_id="active-sessions" → Step 3: command="restore_backup", target_id="db-server"',
+        "expert-apt":          'Step 1: command="detect_c2", target_id="c2-beacon" → Step 2: command="block_outbound", target_id="outbound-fw" → Step 3: command="isolate_host", target_id="infected-host" → Step 4: command="patch_vulnerability", target_id="vuln-cve-2024"',
+    }
+
+    prompt = f"""You are a cybersecurity AI agent for Sentinel-Env.
+
 Task: {task_id}
-Terminal: {obs.get('terminal_output','')}
-Inventory: {json.dumps(obs.get('inventory',[]),indent=2)}
+Terminal output: {terminal}
+Current inventory: {inventory}
 
-Rules:
-- easy-lockdown:     command="lockdown",     target_id="s3-vault"
-- medium-access:     command="revoke_admin", target_id="user-dev-01"
-- hard-breach step1: command="block_ip",     target_id="attacker-ip"
-- hard-breach step2: command="close_port",   target_id="web-server"
+Correct action sequence for this task:
+{rules.get(task_id, 'Follow the terminal instructions.')}
 
-Respond ONLY with JSON: {{"command":"...","target_id":"..."}}"""
+Based on the terminal output, choose the NEXT correct action.
+Respond ONLY with a JSON object. No markdown. No explanation.
+Example: {{"command": "lockdown", "target_id": "s3-vault"}}"""
 
     resp = client.chat.completions.create(
         model=MODEL_NAME,
         messages=[{"role": "user", "content": prompt}],
-        max_tokens=64, temperature=0,
+        max_tokens=64,
+        temperature=0,
     )
-    raw = resp.choices[0].message.content.strip().replace("```json","").replace("```","").strip()
+    raw = resp.choices[0].message.content.strip()
+    raw = raw.replace("```json", "").replace("```", "").strip()
     return json.loads(raw)
 
 # ── LLM: Multi-agent actions ──
@@ -218,6 +233,7 @@ def run_warroom_task() -> None:
           f"rewards={rewards_str}", flush=True)
 
 # ── Main ──
+# ── Main — runs ONLY the task in TASK_NAME ──
 def main():
     try:
         requests.get(f"{SPACE_URL}/health", timeout=15).raise_for_status()
@@ -225,8 +241,18 @@ def main():
         print("[END] success=false steps=0 rewards=0.05", flush=True)
         raise SystemExit(1)
 
-    valid = ["easy-lockdown", "medium-access", "hard-breach", "red-vs-blue"]
-    if TASK_NAME not in valid:
+    valid_tasks = [
+        "easy-lockdown",
+        "easy-secrets",
+        "medium-access",
+        "medium-mfa",
+        "hard-breach",
+        "critical-ransomware",
+        "expert-apt",
+        "red-vs-blue",
+    ]
+
+    if TASK_NAME not in valid_tasks:
         print("[END] success=false steps=0 rewards=0.05", flush=True)
         raise SystemExit(1)