sam25kat/securereview
1
1"""2SecureReview Baseline Inference Script3======================================4Runs an LLM-based agent against the SecureReview environment to produce5baseline scores across all three security review tasks.6 7MANDATORY environment variables:8 API_BASE_URL The API endpoint for the LLM (e.g. https://router.huggingface.co/v1)9 MODEL_NAME The model identifier to use for inference10 HF_TOKEN Your Hugging Face API key11"""12 13import os14import re15import sys16import json17import time18import functools19import requests as http_requests20from openai import OpenAI21 22# All print calls flush stdout immediately so the validator can parse23# [START]/[STEP]/[END] markers in real time.24print = functools.partial(print, flush=True)25 26# === Configuration ===27API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")28MODEL_NAME = os.getenv("MODEL_NAME", "meta-llama/Llama-3.1-8B-Instruct")29HF_TOKEN = os.getenv("HF_TOKEN")30ENV_URL = os.getenv("ENV_URL", "http://localhost:7860")31 32if not HF_TOKEN:33 print("WARNING: HF_TOKEN environment variable not set. LLM calls will fail.")34 print("Set it with: export HF_TOKEN='your-huggingface-token'")35 36client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN or "")37 38TASKS = ["dependency_review", "iac_review", "migration_review"]39 40# === System Prompts ===41SYSTEM_PROMPTS = {42 "dependency_review": """You are a security reviewer analyzing dependency files for supply chain risks.43 44Your job is to identify suspicious packages in the dependency file. Look for:451. Typosquatted packages (misspelled names of popular packages, e.g., 'reqeusts' instead of 'requests')462. Hallucinated/non-existent packages (names that don't exist in the package registry)473. Packages with known critical CVEs at the pinned version484. Deprecated packages with known security issues49 50For each issue found, respond with a JSON action. You MUST respond with exactly ONE JSON object per message.51 52To report a finding:53{54 "action_type": "report_finding",55 "finding": {56 "file": "requirements.txt",57 "line": <line_number>,58 "rule_id": "<DEP-001|DEP-002|DEP-003|DEP-004|DEP-007>",59 "severity": "<critical|high|medium|low>",60 "description": "Explain the issue and name the specific package"61 }62}63 64Rule IDs:65- DEP-001: Package does not exist in registry (hallucinated)66- DEP-002: Package name is typosquat of known package67- DEP-003: Package has known critical CVE68- DEP-004: Package has known high-severity CVE69- DEP-007: Package is deprecated70 71To end the review:72{"action_type": "mark_complete"}73 74Respond with ONLY the JSON object, no explanation.""",75 76 "iac_review": """You are a cloud security reviewer analyzing Infrastructure-as-Code configurations.77 78Check for CIS benchmark violations and security misconfigurations in Terraform/Kubernetes files.79 80FIRST: Use request_file_list to see all files, then request_context for any additional files.81 82For each issue found, respond with a JSON action:83{84 "action_type": "report_finding",85 "finding": {86 "file": "<filename>",87 "line": <line_number>,88 "rule_id": "<IAC-001 through IAC-012>",89 "severity": "<critical|high|medium|low>",90 "description": "Explain the misconfiguration, naming the specific resource"91 }92}93 94Rule IDs:95- IAC-001: Public access to storage resource96- IAC-002: Missing encryption at rest97- IAC-003: Missing encryption in transit98- IAC-004: Overly permissive security group (0.0.0.0/0)99- IAC-005: IAM policy with wildcard actions100- IAC-006: Missing logging/monitoring101- IAC-007: Resource in public subnet without justification102- IAC-008: Missing network access control103- IAC-009: Privileged container/execution104- IAC-010: Cross-account access without restrictions105- IAC-011: Missing backup/recovery configuration106- IAC-012: Hardcoded credentials or secrets107 108Other actions:109{"action_type": "request_file_list"}110{"action_type": "request_context", "filename": "<filename>"}111{"action_type": "mark_complete"}112 113Respond with ONLY the JSON object, no explanation.""",114 115 "migration_review": """You are a database migration safety reviewer analyzing SQL migration scripts.116 117CRITICAL: Before analyzing migrations, request context.json and app_context.py to understand:118- Table sizes (determines if operations will lock tables)119- Deployment strategy (rolling = zero-downtime required)120- Which services depend on which columns121 122Check for unsafe migration patterns:123- MIG-001: Adding NOT NULL column without DEFAULT on large table (causes table lock/rewrite)124- MIG-002: Non-concurrent index creation on large table (blocks writes)125- MIG-003: Dropping column still referenced by application code126- MIG-004: Renaming column during zero-downtime deployment127- MIG-005: Type change with implicit cast on large table128- MIG-006: Migration ordering dependency not satisfied129- MIG-007: Missing expand-migrate-contract pattern130- MIG-008: Foreign key on high-write table without supporting index131- MIG-009: Dropping table with active foreign key references132- MIG-010: Lock-heavy operation without timeout133 134For each issue found:135{136 "action_type": "report_finding",137 "finding": {138 "file": "<migration_file.sql>",139 "line": <line_number>,140 "rule_id": "<MIG-001 through MIG-010>",141 "severity": "<critical|high|medium|low>",142 "description": "Explain WHY the operation is unsafe given the production context (mention table size, deployment strategy). Suggest the safe alternative."143 }144}145 146Other actions:147{"action_type": "request_file_list"}148{"action_type": "request_context", "filename": "<filename>"}149{"action_type": "mark_complete"}150 151Respond with ONLY the JSON object, no explanation.""",152}153 154 155def build_prompt(task_id: str, observation: dict) -> str:156 """Build a user prompt from the current observation."""157 ctx = observation["context"]158 prompt_parts = []159 160 prompt_parts.append(f"Task: {ctx['task_description']}")161 prompt_parts.append(f"Difficulty: {ctx['difficulty']}")162 prompt_parts.append(f"Step: {ctx['current_step']}/{ctx['max_steps']}")163 prompt_parts.append("")164 165 # Review checklist166 prompt_parts.append("Review checklist:")167 for item in ctx["review_checklist"]:168 prompt_parts.append(f" - {item}")169 prompt_parts.append("")170 171 # Files in context172 prompt_parts.append("=== Files to Review ===")173 for f in ctx["files"]:174 prompt_parts.append(f"\n--- {f['filename']} ({f['language']}) ---")175 prompt_parts.append(f["content"])176 prompt_parts.append("")177 178 # Available files not yet loaded179 if ctx["available_files"]:180 prompt_parts.append(181 f"Additional files available (use request_context): {', '.join(ctx['available_files'])}"182 )183 prompt_parts.append("")184 185 # Findings so far186 if observation["findings_so_far"]:187 prompt_parts.append(f"Findings submitted so far: {len(observation['findings_so_far'])}")188 for finding in observation["findings_so_far"]:189 prompt_parts.append(190 f" - [{finding['severity']}] {finding['file']}:{finding.get('line', '?')} "191 f"{finding['rule_id']}: {finding['description'][:80]}..."192 )193 prompt_parts.append("")194 195 # Feedback196 if observation.get("feedback"):197 prompt_parts.append(f"Feedback: {observation['feedback']}")198 prompt_parts.append("")199 200 prompt_parts.append(201 "Analyze the files and respond with your next action as a JSON object."202 )203 204 return "\n".join(prompt_parts)205 206 207def parse_action(text: str) -> dict:208 """Parse LLM output into an action dict."""209 # Strip markdown code fences210 text = re.sub(r"```json\s*", "", text)211 text = re.sub(r"```\s*", "", text)212 text = text.strip()213 214 # Try to find JSON object215 try:216 json_match = re.search(r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}", text, re.DOTALL)217 if json_match:218 data = json.loads(json_match.group())219 # Validate action_type220 if "action_type" in data:221 return data222 except json.JSONDecodeError:223 pass224 225 # Fallback: mark_complete to avoid infinite loops226 return {"action_type": "mark_complete"}227 228 229def run_episode(task_id: str, scenario_id: str = None) -> float:230 """Run a single episode and return the final score.231 232 Emits ``[START]``, ``[STEP]``, and ``[END]`` markers on stdout for233 the validator to parse.234 """235 # === [START] marker ===236 print(f"[START] task={task_id}")237 238 reset_body = {"task_id": task_id}239 if scenario_id:240 reset_body["scenario_id"] = scenario_id241 242 resp = http_requests.post(f"{ENV_URL}/reset", json=reset_body, timeout=30)243 resp.raise_for_status()244 reset_data = resp.json()245 observation = reset_data["observation"]246 info = reset_data["info"]247 scenario = info.get("scenario_id", "unknown")248 249 print(f" Scenario: {scenario}")250 251 done = False252 final_score = 0.0253 step_count = 0254 255 # For migration tasks, start by requesting context files256 first_actions = []257 if task_id == "migration_review":258 available = observation["context"]["available_files"]259 for fname in available:260 if fname in ("context.json", "app_context.py", "service_dependencies.txt"):261 first_actions.append(262 {"action_type": "request_context", "filename": fname}263 )264 elif task_id == "iac_review":265 # Request additional files266 available = observation["context"]["available_files"]267 for fname in available:268 first_actions.append(269 {"action_type": "request_context", "filename": fname}270 )271 272 # Execute pre-planned context requests273 for pre_action in first_actions:274 resp = http_requests.post(275 f"{ENV_URL}/step", json={"action": pre_action}, timeout=30276 )277 resp.raise_for_status()278 step_data = resp.json()279 observation = step_data["observation"]280 done = step_data["done"]281 step_count += 1282 reward_val = step_data.get("reward", 0.0) or 0.0283 final_score = reward_val284 print(f"[STEP] step={step_count} reward={reward_val}")285 if done:286 break287 288 # Main agent loop289 while not done:290 prompt = build_prompt(task_id, observation)291 292 try:293 response = client.chat.completions.create(294 model=MODEL_NAME,295 messages=[296 {"role": "system", "content": SYSTEM_PROMPTS[task_id]},297 {"role": "user", "content": prompt},298 ],299 temperature=0.1,300 max_tokens=500,301 )302 llm_output = response.choices[0].message.content or ""303 except Exception as e:304 print(f" LLM error: {e}")305 llm_output = '{"action_type": "mark_complete"}'306 307 action = parse_action(llm_output)308 309 try:310 resp = http_requests.post(311 f"{ENV_URL}/step", json={"action": action}, timeout=30312 )313 resp.raise_for_status()314 step_data = resp.json()315 except Exception as e:316 print(f" Step error: {e}")317 # Try mark_complete as fallback318 resp = http_requests.post(319 f"{ENV_URL}/step",320 json={"action": {"action_type": "mark_complete"}},321 timeout=30,322 )323 resp.raise_for_status()324 step_data = resp.json()325 326 observation = step_data["observation"]327 done = step_data["done"]328 reward_val = step_data.get("reward", 0.0) or 0.0329 final_score = reward_val330 step_count += 1331 print(f"[STEP] step={step_count} reward={reward_val}")332 333 # Small delay to avoid rate limiting334 time.sleep(0.3)335 336 # === [END] marker ===337 print(f"[END] task={task_id} score={final_score} steps={step_count}")338 return final_score339 340 341def main():342 print("=" * 60)343 print("SecureReview Baseline Inference")344 print("=" * 60)345 print(f"Model: {MODEL_NAME}")346 print(f"API: {API_BASE_URL}")347 print(f"Environment: {ENV_URL}")348 print()349 350 # Get available tasks and scenarios351 tasks_resp = http_requests.get(f"{ENV_URL}/tasks", timeout=10)352 tasks_resp.raise_for_status()353 tasks = tasks_resp.json()354 print(f"Available tasks: {[t['id'] for t in tasks]}")355 print()356 357 all_scores = {}358 359 for task_id in TASKS:360 print(f"\n{'='*40}")361 print(f"Task: {task_id}")362 print(f"{'='*40}")363 364 scores = []365 # Run one episode per task (random scenario)366 score = run_episode(task_id)367 scores.append(score)368 369 avg_score = sum(scores) / len(scores)370 all_scores[task_id] = avg_score371 print(f"\n Average score for {task_id}: {avg_score:.3f}")372 373 # Summary374 print(f"\n{'='*60}")375 print("BASELINE RESULTS SUMMARY")376 print(f"{'='*60}")377 for task_id, score in all_scores.items():378 difficulty = {"dependency_review": "easy", "iac_review": "medium", "migration_review": "hard"}379 print(f" {task_id} ({difficulty[task_id]}): {score:.3f}")380 381 overall = sum(all_scores.values()) / len(all_scores)382 print(f"\n Overall average: {overall:.3f}")383 print(f"{'='*60}")384 385 386if __name__ == "__main__":387 main()388