Vishnunetaran/Sentinel-GRC-Audit
0
1"""2inference.py — GRC Compliance Audit Environment Baseline Inference3==================================================================4MANDATORY (OpenEnv Hackathon Spec):5 API_BASE_URL The API endpoint for the LLM.6 MODEL_NAME The model identifier to use.7 HF_TOKEN Your HuggingFace / OpenAI API key.8 9Optional:10 ENV_BASE_URL Running environment server URL.11 LAUNCH_SERVER Set to '1' to auto-launch the GRC server.12"""13 14from __future__ import annotations15 16import asyncio17import json18import logging19import os20import re21import subprocess22import sys23import time24import traceback25from typing import Any, Dict, List, Optional, Tuple26 27# We catch import errors just in case28try:29 from openai import OpenAI30except ImportError:31 OpenAI = None32 33# ─── Hackathon-required variable names ───────────────────────────────────────34API_BASE_URL: Optional[str] = os.getenv("API_BASE_URL", "https://api.openai.com/v1")35API_KEY: Optional[str] = os.getenv("HF_TOKEN") or os.getenv("API_KEY")36MODEL_NAME: Optional[str] = os.getenv("MODEL_NAME", "gpt-4o-mini")37 38# DO NOT CRASH! Just warn if missing.39missing_vars = []40if not API_BASE_URL: missing_vars.append("API_BASE_URL")41if not API_KEY: missing_vars.append("HF_TOKEN / API_KEY")42if not MODEL_NAME: missing_vars.append("MODEL_NAME")43if missing_vars:44 print(f"WARNING: Missing required environment variables: {', '.join(missing_vars)}")45 46# ─── GRC-specific config ─────────────────────────────────────────────────────47ENV_BASE_URL: str = os.getenv("ENV_BASE_URL", "http://localhost:8000")48LAUNCH_SERVER: bool = os.getenv("LAUNCH_SERVER", "1") == "1" # Auto-start server by default49STEP_DELAY: float = float(os.getenv("STEP_DELAY", "3")) # Seconds between LLM calls50TEMPERATURE: float = 0.051MAX_TOKENS: int = 120052DEBUG: bool = os.getenv("DEBUG", "0") == "1"53 54TASK_IDS = ["task_easy", "task_medium", "task_hard"]55MAX_STEPS = {"task_easy": 5, "task_medium": 10, "task_hard": 20}56 57logging.basicConfig(58 level=logging.DEBUG if DEBUG else logging.INFO,59 format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",60)61logger = logging.getLogger("grc_inference")62logging.getLogger("httpx").setLevel(logging.WARNING)63 64def safe_print(line: str) -> None:65 sys.stdout.write(line + "\n")66 sys.stdout.flush()67 time.sleep(0.05)68 69def safe_score(x: float) -> float:70 x = float(x)71 if x <= 0.0: return 0.000172 if x >= 1.0: return 0.999973 return x74 75# ─────────────────────────────────────────────────────────────────────────────76# SYSTEM PROMPT — High-Fidelity GRC Auditor77# ─────────────────────────────────────────────────────────────────────────────78 79_cached_system_prompt: Optional[str] = None80 81def get_system_prompt() -> str:82 global _cached_system_prompt83 if _cached_system_prompt is not None:84 return _cached_system_prompt85 86 base_dir = os.path.dirname(os.path.abspath(__file__))87 tax_dir = os.path.join(base_dir, "grc_compliance_audit_env", "server", "data", "taxonomies")88 89 iso_lines, nist_lines, soc2_lines = [], [], []90 91 try:92 with open(os.path.join(tax_dir, "iso27001_controls.json"), "r") as f:93 t = json.load(f)94 for c in t.get("controls", []):95 if c["id"].startswith("A.5") or c["id"].startswith("A.6") or c["id"].startswith("A.7") or c["id"].startswith("A.8"):96 iso_lines.append(f" {c['id']}: {c['name']}")97 except Exception as e:98 logger.warning(f"Failed to load ISO taxonomy: {e}")99 100 try:101 with open(os.path.join(tax_dir, "nist_80053_families.json"), "r") as f:102 t = json.load(f)103 for fam in t.get("families", []):104 for c in fam.get("key_controls", []):105 nist_lines.append(f" {c['id']}: {c['name']}")106 except Exception as e:107 logger.warning(f"Failed to load NIST taxonomy: {e}")108 109 try:110 with open(os.path.join(tax_dir, "soc2_tsc.json"), "r") as f:111 t = json.load(f)112 for cat in t.get("categories", []):113 for c in cat.get("criteria", []):114 soc2_lines.append(f" {c['id']}: {c['name']}")115 except Exception as e:116 logger.warning(f"Failed to load SOC2 taxonomy: {e}")117 118 prompt = f"""You are a certified GRC (Governance, Risk, and Compliance) analyst.119 120Your task is to audit a policy document against one or more compliance frameworks and return a structured JSON response.121 122════════════════════════════════════════════════════════════════123CONTROL ID REFERENCE (STRICT ALPHANUMERIC ENFORCEMENT)124════════════════════════════════════════════════════════════════125The following are the ONLY valid Control IDs you may use. You must map the document to these EXACT IDs based on their descriptions.126 127--- ISO 27001:2022 ---128{chr(10).join(iso_lines)}129 130--- NIST SP 800-53 Rev 5 ---131{chr(10).join(nist_lines)}132 133--- SOC 2 Trust Services Criteria ---134{chr(10).join(soc2_lines)}135 136════════════════════════════════════════════════════════════════137RULES & REQUIREMENTS (CRITICAL)138════════════════════════════════════════════════════════════════1391. You must output ONLY the alphanumeric Control IDs (e.g., A.5.15, AC-2, CC6.1). Do not use descriptive text or names in the mapping lists. If you use a name instead of an ID, the score will be 0.1402. Output ONLY raw JSON. No preamble, no 'Here is your audit', no markdown fences, no explanation. The very FIRST character must be {{ and the last must be }}.1413. BE EXHAUSTIVE but PRECISE: Map every applicable control, but only include IDs whose description closely matches the section text. Do not guess. Each section typically maps to 1-3 controls per framework.1424. In Task 3, look for sections that satisfy ISO, NIST, and SOC 2 simultaneously to trigger the +0.05 Alignment Bonus ★.1435. risk_level MUST be exactly one of: critical, high, medium, low.1446. framework MUST be exactly one of: iso27001, nist_80053, soc2.1457. A shared_control entry MUST always contain ALL THREE fields: iso_control_id, nist_control_id, AND soc2_criteria_id. If any of the three is missing or unknown, do NOT include that entry in shared_controls — put the partial mapping in control_mappings instead.1468. For gaps: look for controls that the applicable frameworks REQUIRE but that the policy text does NOT address.147 148OUTPUT FORMAT:149{{150 "task_id": "<task_id>",151 "reasoning": "Brief analysis",152 "control_mappings": [153 {{154 "section_id": "section_1",155 "iso_control_ids": ["A.5.15", "A.5.16"],156 "nist_control_ids": ["AC-2", "AC-3"],157 "soc2_criteria_ids": ["CC6.1", "CC6.2"]158 }}159 ],160 "gaps": [161 {{162 "control_id": "A.8.5",163 "framework": "iso27001",164 "risk_level": "critical",165 "gap_description": "MFA is required but not mentioned.",166 "affected_section": "section_1",167 "remediation": "Implement MFA"168 }}169 ],170 "shared_controls": [171 {{172 "policy_section_id": "section_1",173 "iso_control_id": "A.5.15",174 "nist_control_id": "AC-2",175 "soc2_criteria_id": "CC6.1"176 }}177 ],178 "executive_summary": "Overall summary"179}}180"""181 _cached_system_prompt = prompt182 return prompt183 184# ─────────────────────────────────────────────────────────────────────────────185# Server management 186# ─────────────────────────────────────────────────────────────────────────────187 188_server_proc: Optional[subprocess.Popen] = None189 190def start_server_local() -> None:191 global _server_proc192 try:193 pkg_dir = os.path.dirname(os.path.abspath(__file__))194 cmd = [195 sys.executable, "-m", "uvicorn",196 "grc_compliance_audit_env.server.app:app",197 "--host", "0.0.0.0", "--port", "8000",198 "--log-level", "warning",199 ]200 logger.info("Starting GRC server (from root directory)...")201 _server_proc = subprocess.Popen(cmd, cwd=pkg_dir)202 import urllib.request203 for attempt in range(30):204 time.sleep(1)205 try:206 urllib.request.urlopen(f"http://localhost:8000/health", timeout=2)207 logger.info("Server ready.")208 return209 except Exception:210 pass211 logger.warning("GRC server failed to start, but continuing execution safely.")212 except Exception as e:213 logger.error(f"Error starting local server: {e}")214 215def stop_server_local() -> None:216 global _server_proc217 if _server_proc is not None:218 try:219 _server_proc.terminate()220 _server_proc.wait(timeout=5)221 _server_proc = None222 logger.info("GRC server stopped.")223 except Exception as e:224 logger.error(f"Error stopping local server: {e}")225 226# ─────────────────────────────────────────────────────────────────────────────227# LLM Logic228# ─────────────────────────────────────────────────────────────────────────────229 230_client: Optional[Any] = None231 232def get_openai_client() -> Optional[Any]:233 global _client234 if _client is None:235 try:236 if not OpenAI:237 logger.error("OpenAI library not found. LLM unavailable.")238 return None239 if not API_BASE_URL or not API_KEY:240 logger.error("Missing API URL or KEY. LLM unavailable.")241 return None242 _client = OpenAI(243 base_url=API_BASE_URL,244 api_key=API_KEY,245 max_retries=0, 246 timeout=30.0 # Strict timeout bound247 )248 except Exception as e:249 logger.error(f"Failed to cleanly initialize OpenAI client: {e}")250 return None251 return _client252 253def build_user_prompt(254 obs: Dict[str, Any], 255 target_section: Optional[int], 256 prev_feedback: str, 257 cumulative_action: Dict[str, Any]258) -> str:259 task_id = obs.get("task_id", "")260 frameworks = obs.get("target_frameworks", [])261 total_sections = obs.get("total_sections", 1)262 263 lines = [264 f"TASK: {obs.get('task_description', '')}",265 f"TARGET FRAMEWORKS: {', '.join(frameworks)}",266 ]267 268 if prev_feedback and "Episode started" not in prev_feedback:269 lines.append(f"\n[GRADER FEEDBACK]\n{prev_feedback}")270 271 lines.append(f"\n{'─' * 60}\nPOLICY TEXT:\n{'─' * 60}")272 lines.append(obs.get("policy_text", ""))273 lines.append(f"{'─' * 60}")274 275 if target_section is not None:276 lines.append(f"\n[PROGRESSIVE AUDIT MODE]")277 lines.append(f"In this step, audit ONLY 'section_{target_section}'.")278 else:279 lines.append(f"\n[REFINEMENT MODE]")280 lines.append("You have audited all sections. Review the GRADER FEEDBACK.")281 lines.append("CRITICAL: DO NOT RE-OUTPUT ANY CONTROLS OR GAPS THAT ARE ALREADY IN 'MEMORY'.")282 283 known_maps = [f"{m.get('section_id')}: ISO={m.get('iso_control_ids',[])} NIST={m.get('nist_control_ids',[])} SOC2={m.get('soc2_criteria_ids',[])}" for m in cumulative_action.get("control_mappings", [])]284 lines.append(f"Mappings already submitted: {known_maps}")285 return "\n".join(lines)286 287def call_llm(user_prompt: str, task_id: str) -> Dict[str, Any]:288 # SAFE FALLBACK NOOP ACTION289 noop_action = {290 "task_id": task_id, 291 "control_mappings": [], 292 "gaps": [], 293 "shared_controls": []294 }295 296 try:297 llm = get_openai_client()298 if not llm:299 return noop_action300 301 time.sleep(STEP_DELAY)302 303 for attempt in range(3):304 try:305 response = llm.chat.completions.create(306 model=MODEL_NAME,307 messages=[308 {"role": "system", "content": get_system_prompt()},309 {"role": "user", "content": user_prompt},310 ],311 temperature=TEMPERATURE,312 max_tokens=MAX_TOKENS,313 )314 315 if not response or not response.choices or not response.choices[0].message:316 continue317 318 raw = response.choices[0].message.content or "{}"319 except Exception as api_exc:320 if "429" in str(api_exc) or "RateLimitError" in str(type(api_exc)):321 time.sleep(10)322 continue323 324 if attempt == 2:325 return noop_action326 continue327 328 raw = re.sub(r'^```(?:json)?\s*', '', raw.strip())329 raw = re.sub(r'\s*```$', '', raw)330 331 try:332 action_dict = json.loads(raw)333 if isinstance(action_dict, dict):334 action_dict["task_id"] = task_id335 action_dict.setdefault("control_mappings", [])336 action_dict.setdefault("gaps", [])337 action_dict.setdefault("shared_controls", [])338 return action_dict339 except Exception:340 pass341 342 return noop_action343 344 except Exception as e:345 logger.error(f"Absolute failure in call_llm wrapper: {e}")346 return noop_action347 348# ─────────────────────────────────────────────────────────────────────────────349# Core Async Task Runner350# ─────────────────────────────────────────────────────────────────────────────351 352def merge_actions(cumulative: Dict[str, Any], new_action: Dict[str, Any]) -> None:353 try:354 for new_map in new_action.get("control_mappings", []) or []:355 sec_id = new_map.get("section_id")356 existing = next((m for m in cumulative["control_mappings"] if m.get("section_id") == sec_id), None)357 if existing:358 existing["iso_control_ids"] = list(set(existing.get("iso_control_ids", []) + (new_map.get("iso_control_ids") or [])))359 existing["nist_control_ids"] = list(set(existing.get("nist_control_ids", []) + (new_map.get("nist_control_ids") or [])))360 existing["soc2_criteria_ids"] = list(set(existing.get("soc2_criteria_ids", []) + (new_map.get("soc2_criteria_ids") or [])))361 else:362 cumulative["control_mappings"].append(new_map)363 364 for new_gap in new_action.get("gaps", []) or []:365 dup = any(g.get("control_id") == new_gap.get("control_id") and g.get("affected_section") == new_gap.get("affected_section") for g in cumulative["gaps"])366 if not dup:367 cumulative["gaps"].append(new_gap)368 369 for new_sc in new_action.get("shared_controls", []) or []:370 dup = any(sc.get("policy_section_id") == new_sc.get("policy_section_id") for sc in cumulative["shared_controls"])371 if not dup:372 cumulative["shared_controls"].append(new_sc)373 except Exception as e:374 logger.error(f"merge_actions failed: {e}")375 376 377CYAN, GREEN, YELLOW, RED, BOLD, DIM, RESET = "\033[96m", "\033[92m", "\033[93m", "\033[91m", "\033[1m", "\033[2m", "\033[0m"378 379async def run_task(task_id: str, ws_base_url: str) -> Tuple[float, bool, int]:380 try:381 import websockets as _ws382 ws_url = ws_base_url.replace("http://", "ws://").replace("https://", "wss://") + "/ws"383 max_steps = MAX_STEPS.get(task_id, 10)384 385 print(f"\n{BOLD}{'▶'} {task_id.upper()}{RESET} (max {max_steps} steps)")386 safe_print(f"[START] {task_id}")387 388 cumulative_action = {389 "task_id": task_id,390 "control_mappings": [],391 "gaps": [],392 "shared_controls": []393 }394 395 max_reward_seen = -1.0396 best_step_reward = 0.0397 decay_steps = 0398 final_obs = {}399 step = 0400 401 try:402 async with _ws.connect(ws_url, max_size=20_000_000, ping_timeout=20, ping_interval=20) as ws:403 await ws.send(json.dumps({"type": "reset", "options": {"task_id": task_id}}))404 msg = json.loads(await ws.recv())405 obs = msg.get("data", {})406 407 prev_feedback = obs.get("grader_feedback", "")408 final_obs = obs409 total_sections = obs.get("total_sections", 1)410 411 for step in range(1, max_steps + 1):412 if obs.get("done", False):413 break414 415 safe_print(f"[STEP] {task_id} step={step}")416 417 target_section = step if step <= total_sections else None418 user_prompt = build_user_prompt(obs, target_section, prev_feedback, cumulative_action)419 420 try:421 new_action = await asyncio.to_thread(call_llm, user_prompt, task_id)422 except Exception as exc:423 logger.error("LLM thread crashed: %s", exc)424 break425 426 merge_actions(cumulative_action, new_action)427 428 await ws.send(json.dumps({"type": "step", "action": cumulative_action}))429 msg = json.loads(await ws.recv())430 431 if msg.get("type") == "error":432 logger.error("Server API error: %s", msg.get("message"))433 break434 435 obs = msg.get("data", {})436 final_obs = obs437 438 step_reward = obs.get("step_reward", 0.0)439 best_step_reward = max(best_step_reward, step_reward)440 print(f" Step {step}/{max_steps} | Target: {f'Section {target_section}' if target_section else 'Refinement'} | Step Reward: {step_reward:+.3f} (best: {best_step_reward:.3f})")441 442 if target_section is None:443 if step_reward < max_reward_seen: decay_steps += 1444 else: decay_steps = 0445 max_reward_seen = max(max_reward_seen, step_reward)446 if decay_steps >= 2: break447 else:448 max_reward_seen = max(max_reward_seen, step_reward)449 450 prev_feedback = obs.get("grader_feedback", "")451 452 episode_done = final_obs.get("done", False) or (step > 0)453 return best_step_reward, episode_done, step454 455 except Exception as e:456 logger.error(f"WebSocket execution error for task {task_id}: {e}")457 return best_step_reward, False, step458 finally:459 best_step_reward = safe_score(best_step_reward)460 safe_print(f"[END] {task_id} score={best_step_reward:.4f} steps={step}")461 462 except Exception as e:463 logger.error(f"Fatal error preparing task {task_id}: {e}")464 safe_print(f"[END] {task_id} score=0.0001 steps=0")465 return 0.0001, False, 0466 467async def main_async() -> None:468 print(f"{BOLD} GRC Compliance Audit — Progressive Inference{RESET}")469 470 try:471 if LAUNCH_SERVER: start_server_local()472 except Exception:473 pass474 475 results = []476 477 for task_id in TASK_IDS:478 try:479 score, done, steps = await run_task(task_id, ENV_BASE_URL)480 results.append({"id": task_id, "score": score, "done": done, "steps": steps})481 except Exception as exc:482 logger.error(f"Caught unhandled task error: {exc}")483 results.append({"id": task_id, "score": 0.0001, "done": False, "steps": 0, "err": str(exc)})484 485 try:486 if LAUNCH_SERVER: stop_server_local()487 except Exception:488 pass489 490 print(f"\n{BOLD} FINAL SCOREBOARD{RESET}")491 total = 0.0492 for r in results:493 sym = f"{GREEN}✓{RESET}" if r["score"] > 0.01 else f"{RED}✗{RESET}"494 score = r["score"]495 col = GREEN if score >= 0.6 else YELLOW if score >= 0.3 else RED496 print(f" {r['id']:<15} {col}{score:.4f}{RESET} {sym} {r['steps']} steps")497 total += score498 499 print(f" AVERAGE: {GREEN if total/3 >= 0.5 else YELLOW}{(total/3):.4f}{RESET}\n")500 501if __name__ == "__main__":502 try:503 asyncio.run(main_async())504 except Exception as e:505 print(f"FATAL UNHANDLED EXCEPTION PREVENTED CRASH: {e}")506 print("\n FINAL SCOREBOARD")507 print(" task_easy 0.0001 ✗ 0 steps")508 print(" task_medium 0.0001 ✗ 0 steps")509 print(" task_hard 0.0001 ✗ 0 steps")510 print(" AVERAGE: 0.0001")511 sys.exit(0) # IMPORTANT: Exit cleanly512