Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| """ | |
| inference.py — Sentinel Env Multi-Agent Inference Script | |
| OpenEnv 2026 Architecture — Session-isolated, WebSocket-ready | |
| Follows [START]/[STEP]/[END] stdout format. | |
| Runs ONE task per invocation via TASK_NAME env var. | |
| """ | |
| import os | |
| import json | |
| import requests | |
| from openai import OpenAI | |
| # ── Credentials ── | |
| API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1") | |
| MODEL_NAME = os.getenv("MODEL_NAME", "meta-llama/Meta-Llama-3-8B-Instruct") | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| if HF_TOKEN is None: | |
| raise ValueError("HF_TOKEN environment variable is required") | |
| SPACE_URL = os.getenv("SPACE_URL", "https://rudrapatel-1908-sentinel-env.hf.space") | |
| TASK_NAME = os.getenv("TASK_NAME", "easy-lockdown") | |
| BENCHMARK = "sentinel_env" | |
| MAX_STEPS = 8 | |
| client = OpenAI(api_key=HF_TOKEN, base_url=API_BASE_URL) | |
| # ── HTTP Helpers ── | |
| def env_reset(task_id: str) -> dict: | |
| r = requests.post(f"{SPACE_URL}/reset", json={"task_id": task_id}, timeout=30) | |
| r.raise_for_status() | |
| return r.json() | |
| def env_step(command: str, target_id: str) -> dict: | |
| r = requests.post(f"{SPACE_URL}/step", | |
| json={"command": command, "target_id": target_id}, timeout=30) | |
| r.raise_for_status() | |
| return r.json() | |
| def warroom_reset() -> dict: | |
| r = requests.post(f"{SPACE_URL}/warroom/reset", | |
| json={"task_id": "red-vs-blue"}, timeout=30) | |
| r.raise_for_status() | |
| return r.json() | |
| def warroom_step(att_type, att_res, scn_type, scn_tip, scn_res, rem_type, rem_tip, rem_res) -> dict: | |
| payload = { | |
| "attacker": {"action_type": att_type, "target_resource": att_res}, | |
| "scanner": {"action_type": scn_type, "target_ip": scn_tip, "target_resource": scn_res}, | |
| "remediator": {"action_type": rem_type, "target_ip": rem_tip, "target_resource": rem_res}, | |
| } | |
| r = requests.post(f"{SPACE_URL}/warroom/step", json=payload, timeout=30) | |
| r.raise_for_status() | |
| return r.json() | |
| # ── LLM: Single-agent action ── | |
| def get_single_action(task_id: str, obs: dict) -> dict: | |
| 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 output: {terminal} | |
| Current inventory: {inventory} | |
| Correct action sequence for this task: | |
| {rules.get(task_id, 'Follow the terminal instructions.')} | |
| 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, | |
| ) | |
| raw = resp.choices[0].message.content.strip() | |
| raw = raw.replace("```json", "").replace("```", "").strip() | |
| return json.loads(raw) | |
| # ── LLM: Multi-agent actions ── | |
| def _ask_llm(prompt: str) -> dict: | |
| try: | |
| resp = client.chat.completions.create( | |
| model=MODEL_NAME, | |
| messages=[{"role": "user", "content": prompt}], | |
| max_tokens=80, temperature=0, | |
| ) | |
| raw = resp.choices[0].message.content.strip().replace("```json","").replace("```","").strip() | |
| return json.loads(raw) | |
| except Exception: | |
| return {} | |
| def get_multi_actions(obs: dict) -> dict: | |
| scn_obs = obs.get("scanner_obs", {}) | |
| rem_obs = obs.get("remediator_obs", {}) | |
| att_obs = obs.get("attacker_obs", {}) | |
| att = _ask_llm(f"""You are the Red Team attacker. | |
| Visible resources: {att_obs.get('visible_resources',[])} | |
| Position: {att_obs.get('current_position','unknown')} | |
| Data stolen: {att_obs.get('data_exfiltrated',False)} | |
| action_type options: port_scan, privilege_escalation, lateral_movement, data_exfiltration, hide_tracks | |
| target_resource options: s3-vault, db-server, secret-vault, web-server, user-admin | |
| Respond ONLY with JSON: {{"action_type":"...","target_resource":"..."}}""") | |
| scn = _ask_llm(f"""You are the Blue Team Scanner. | |
| IP connections: {scn_obs.get('ip_connections',[])} | |
| Flagged threats: {scn_obs.get('flagged_threats',[])} | |
| Logs: {scn_obs.get('security_logs',[])[-2:]} | |
| action_type options: scan_logs, flag_threat, share_intel | |
| Respond ONLY with JSON: {{"action_type":"...","target_ip":"45.33.32.156","target_resource":"web-server"}}""") | |
| rem = _ask_llm(f"""You are the Blue Team Remediator. | |
| Open ports: {rem_obs.get('open_ports',[])} | |
| Locked resources: {rem_obs.get('locked_resources',[])} | |
| Intel from Scanner: {rem_obs.get('received_intel',[])} | |
| action_type options: block_ip, close_port, lockdown, revoke_access | |
| Respond ONLY with JSON: {{"action_type":"...","target_ip":"45.33.32.156","target_resource":"web-server"}}""") | |
| # Safe defaults | |
| if not att.get("action_type"): att = {"action_type": "port_scan", "target_resource": "web-server"} | |
| if not scn.get("action_type"): scn = {"action_type": "scan_logs", "target_resource": "web-server"} | |
| if not rem.get("action_type"): rem = {"action_type": "close_port", "target_resource": "web-server"} | |
| return {"attacker": att, "scanner": scn, "remediator": rem} | |
| # ── Single-agent runner ── | |
| def run_single_task(task_id: str) -> None: | |
| result = env_reset(task_id) | |
| obs = result["observation"] | |
| rewards = [] | |
| done = False | |
| success = False | |
| print(f"[START] task={task_id} env={BENCHMARK} model={MODEL_NAME}", flush=True) | |
| for step_num in range(1, MAX_STEPS + 1): | |
| action_str = "null" | |
| reward = 0.05 | |
| error = "null" | |
| try: | |
| action = get_single_action(task_id, obs) | |
| cmd = action["command"] | |
| tgt = action["target_id"] | |
| action_str = f"{cmd}({tgt})" | |
| result = env_step(cmd, tgt) | |
| obs = result["observation"] | |
| reward = float(result["reward"]) | |
| done = result["done"] | |
| except Exception as e: | |
| error = str(e).replace("\n", " ") | |
| rewards.append(reward) | |
| print(f"[STEP] step={step_num} action={action_str} " | |
| f"reward={reward:.2f} done={str(done).lower()} error={error}", flush=True) | |
| if done: | |
| success = True | |
| break | |
| rewards_str = ",".join(f"{r:.2f}" for r in rewards) | |
| print(f"[END] success={str(success).lower()} steps={len(rewards)} " | |
| f"rewards={rewards_str}", flush=True) | |
| # ── Multi-agent runner ── | |
| def run_warroom_task() -> None: | |
| result = warroom_reset() | |
| obs = result | |
| rewards = [] | |
| done = False | |
| success = False | |
| print(f"[START] task=red-vs-blue env={BENCHMARK} model={MODEL_NAME}", flush=True) | |
| for step_num in range(1, MAX_STEPS + 1): | |
| action_str = "null" | |
| reward = 0.05 | |
| error = "null" | |
| try: | |
| acts = get_multi_actions(obs) | |
| att = acts["attacker"] | |
| scn = acts["scanner"] | |
| rem = acts["remediator"] | |
| action_str = ( | |
| f"red:{att.get('action_type','?')}({att.get('target_resource','?')})|" | |
| f"scan:{scn.get('action_type','?')}|" | |
| f"rem:{rem.get('action_type','?')}" | |
| ) | |
| result = warroom_step( | |
| att.get("action_type","port_scan"), att.get("target_resource","web-server"), | |
| scn.get("action_type","scan_logs"), scn.get("target_ip",""), scn.get("target_resource","web-server"), | |
| rem.get("action_type","close_port"), rem.get("target_ip",""), rem.get("target_resource","web-server"), | |
| ) | |
| obs = result.get("observation", obs) | |
| rwds = result.get("rewards", {}) | |
| done = result.get("done", False) | |
| blue = (rwds.get("scanner", 0.05) + rwds.get("remediator", 0.05)) / 2 | |
| reward = round(max(0.01, min(0.99, blue)), 2) | |
| except Exception as e: | |
| error = str(e).replace("\n", " ") | |
| rewards.append(reward) | |
| print(f"[STEP] step={step_num} action={action_str} " | |
| f"reward={reward:.2f} done={str(done).lower()} error={error}", flush=True) | |
| if done: | |
| att_obs = obs.get("attacker_obs", {}) if isinstance(obs, dict) else {} | |
| success = not att_obs.get("data_exfiltrated", True) | |
| break | |
| rewards_str = ",".join(f"{r:.2f}" for r in rewards) | |
| print(f"[END] success={str(success).lower()} steps={len(rewards)} " | |
| 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() | |
| except Exception: | |
| print("[END] success=false steps=0 rewards=0.05", flush=True) | |
| raise SystemExit(1) | |
| 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) | |
| try: | |
| if TASK_NAME == "red-vs-blue": | |
| run_warroom_task() | |
| else: | |
| run_single_task(TASK_NAME) | |
| except Exception: | |
| print("[END] success=false steps=0 rewards=0.05", flush=True) | |
| raise SystemExit(1) | |
| if __name__ == "__main__": | |
| main() |