bothari01/secops-env
0
1#!/usr/bin/env python32"""3SecOps Environment - Inference Script4 5MANDATORY:6- Before submitting, ensure the following variables are defined:7 API_BASE_URL The API endpoint for the LLM.8 MODEL_NAME The model identifier to use for inference.9 HF_TOKEN Your Hugging Face / API key.10 11- The inference script must be named `inference.py` and placed in the root directory12- Participants must use OpenAI Client for all LLM calls using above variables13 14STDOUT FORMAT:15 [START] task=<task_name> env=<benchmark> model=<model_name>16 [STEP] step=<n> action=<action_str> reward=<0.00> done=<true|false> error=<msg|null>17 [END] success=<true|false> steps=<n> score=<score> rewards=<r1,r2,...,rn>18"""19 20import os21import sys22import json23import re24import textwrap25from typing import List, Dict, Any, Optional26from dataclasses import dataclass, field27from openai import OpenAI28 29from secops_env import SecOpsEnv, SecOpsAction30from secops_env.models import TaskType, ActionType31 32 33API_BASE_URL = os.getenv("API_BASE_URL", "https://router.huggingface.co/v1")34MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-7B-Instruct")35HF_TOKEN = os.getenv("HF_TOKEN")36TASK_NAME = os.getenv("TASK_NAME", "pii_redaction")37BENCHMARK = os.getenv("BENCHMARK", "secops_env")38MAX_STEPS = int(os.getenv("MAX_STEPS", "10"))39TEMPERATURE = float(os.getenv("TEMPERATURE", "0.7"))40MAX_TOKENS = int(os.getenv("MAX_TOKENS", "500"))41SUCCESS_SCORE_THRESHOLD = 0.142# #0.01 = 1e-943DEBUG = os.getenv("DEBUG", "false").lower() == "true"44 45 46ALL_TASKS = [47 ("pii_redaction", "easy"),48 ("public_access", "medium"),49 ("ghost_user", "hard"),50 ("log_analysis", "medium"),51 ("config_hardening", "hard"),52]53 54 55def log_start(task: str, env: str, model: str) -> None:56 print(f"[START] task={task} env={env} model={model}", flush=True)57 58 59def log_step(60 step: int, action: str, reward: float, done: bool, error: Optional[str]61) -> None:62 error_val = error if error else "null"63 done_val = str(done).lower()64 action_str = action.replace('"', '\\"')[:100] if action else "noop"65 print(66 f'[STEP] step={step} action="{action_str}" reward={reward:.2f} done={done_val} error={error_val}',67 flush=True,68 )69 70 71def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:72 rewards_str = ",".join(f"{r:.2f}" for r in rewards)73 print(74 f"[END] success={str(success).lower()} steps={steps} score={score:.3f} rewards={rewards_str}",75 flush=True,76 )77 78 79@dataclass80class AgentConfig:81 model_name: str = MODEL_NAME82 api_base_url: str = API_BASE_URL83 api_key: Optional[str] = HF_TOKEN84 max_steps: int = MAX_STEPS85 temperature: float = TEMPERATURE86 max_tokens: int = MAX_TOKENS87 88 89class SecOpsAgent:90 TASK_PROMPTS = {91 "pii_redaction": textwrap.dedent("""92 You are a security analyst specializing in PII redaction.93 94 Your task: Identify and redact all Personally Identifiable Information (PII) from the given text.95 96 PII types to look for:97 - SSN: XXX-XX-XXXX format98 - Email: user@domain.com format99 - Phone: XXX-XXX-XXXX or +1-XXX-XXX-XXXX format100 - Credit Card: XXXX-XXXX-XXXX-XXXX format101 - IP Address: XXX.XXX.XXX.XXX format102 103 IMPORTANT:104 1. Replace ALL PII with [REDACTED]105 2. Preserve the original text structure106 3. Do not modify non-PII content107 4. Return ONLY the redacted text108 109 Respond with your redacted version.110 """).strip(),111 "public_access": textwrap.dedent("""112 You are a cloud security analyst.113 114 Your task: Identify S3 buckets with public access and create a plan to fix them.115 116 Given a list of cloud resources, identify which ones have public access enabled.117 118 Respond in JSON format:119 {"identified_public": ["bucket-name-1", "bucket-name-2"], "reasoning": "Brief explanation"}120 121 If no public buckets, respond:122 {"identified_public": [], "reasoning": "No public buckets found"}123 """).strip(),124 "ghost_user": textwrap.dedent("""125 You are a security operations analyst.126 127 Your task: Identify orphaned/inactive ghost user accounts that should be disabled.128 129 Ghost user criteria:130 - No login in 90+ days131 - No active cloud resources132 - No recent deployments133 134 Given user account data, identify ghost users.135 136 Respond in JSON format:137 {"identified_ghosts": ["user@company.com", "..."], "reasoning": "Brief explanation"}138 139 If no ghost users, respond:140 {"identified_ghosts": [], "reasoning": "No ghost users found"}141 """).strip(),142 "log_analysis": textwrap.dedent("""143 You are a security analyst specializing in SIEM and log analysis.144 145 Your task: Analyze the provided security logs and classify the security event.146 147 Classifications: MALWARE, TRUE_POSITIVE, FALSE_POSITIVE, NEEDS_INVESTIGATION,148 LATERAL_MOVEMENT, DATA_EXFILTRATION, UNAUTHORIZED_ACCESS, BENIGN149 150 Severity levels: LOW, MEDIUM, HIGH, CRITICAL151 152 Respond in JSON format:153 {"classification": "MALWARE", "severity": "HIGH", "reasoning": "Brief explanation"}154 """).strip(),155 "config_hardening": textwrap.dedent("""156 You are a cloud security engineer specializing in configuration review.157 158 Your task: Review the provided configuration for security issues.159 160 Common issues: privileged containers, running as root, overly permissive IAM,161 plaintext secrets, public S3 access, missing TLS, exposed services.162 163 Respond in JSON format:164 {"config_issues": [{"type": "issue_type", "severity": "HIGH", "fix": "description"}],165 "hardened_config": "Full corrected configuration"}166 """).strip(),167 }168 169 def __init__(self, config: AgentConfig):170 self.config = config171 self.client = None172 if config.api_key:173 self.client = OpenAI(base_url=config.api_base_url, api_key=config.api_key)174 175 def build_prompt(self, observation) -> str:176 task_type = observation.task_type177 if hasattr(task_type, "value"):178 task_type = task_type.value179 180 context = observation.context181 prompt_parts = [self.TASK_PROMPTS.get(task_type, "Complete the security task.")]182 prompt_parts.append(f"\n\nObjective: {observation.objective}")183 184 if task_type == "pii_redaction":185 text = context.get("text", "")186 prompt_parts.append(f"\n\nText to redact:\n{text}")187 elif task_type == "public_access":188 resources = context.get("resources", [])189 prompt_parts.append("\n\nCloud Resources:")190 for r in resources:191 pub_status = "PUBLIC" if r.get("public") else "private"192 prompt_parts.append(f" - {r['name']} ({pub_status})")193 elif task_type == "ghost_user":194 users = context.get("users", [])195 prompt_parts.append("\n\nUser Accounts:")196 for u in users:197 days = u.get("last_login", "unknown")198 resources = len(u.get("active_resources", []))199 prompt_parts.append(200 f" - {u['username']}: last login {days}, resources: {resources}"201 )202 elif task_type == "log_analysis":203 logs = context.get("logs", "")204 prompt_parts.append(f"\n\nSecurity Logs to Analyze:\n{logs}")205 elif task_type == "config_hardening":206 config_content = context.get("config_content", "")207 config_type = context.get("config_type", "yaml")208 prompt_parts.append(209 f"\n\nConfiguration to Review ({config_type}):\n{config_content}"210 )211 212 if observation.feedback:213 prompt_parts.append(f"\n\nPrevious feedback: {observation.feedback}")214 215 prompt_parts.append(216 f"\n\nStep {observation.step_count + 1}/{observation.max_steps}"217 )218 return "\n".join(prompt_parts)219 220 def _extract_json(self, text: str) -> Optional[Dict[str, Any]]:221 json_pattern = r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}"222 matches = re.findall(json_pattern, text, re.DOTALL)223 for match in matches:224 try:225 return json.loads(match)226 except json.JSONDecodeError:227 continue228 return None229 230 def _extract_redacted_text(self, text: str) -> Optional[str]:231 lines = text.strip().split("\n")232 for line in lines:233 if "[REDACTED]" in line or "[redacted]" in line.lower():234 return line.strip()235 if "```" in text:236 parts = text.split("```")237 for part in parts:238 if "[REDACTED]" in part:239 return part.strip()240 return None241 242 def generate_action(self, observation) -> SecOpsAction:243 task_type = observation.task_type244 if hasattr(task_type, "value"):245 task_type = task_type.value246 247 if not self.client:248 return self._fallback_action(observation, task_type)249 250 prompt = self.build_prompt(observation)251 252 try:253 completion = self.client.chat.completions.create(254 model=self.config.model_name,255 messages=[256 {257 "role": "system",258 "content": "You are a security operations assistant.",259 },260 {"role": "user", "content": prompt},261 ],262 temperature=self.config.temperature,263 max_tokens=self.config.max_tokens,264 )265 response_text = completion.choices[0].message.content or ""266 except Exception as e:267 if DEBUG:268 print(f"[DEBUG] API error: {type(e).__name__}: {e}", flush=True)269 return self._fallback_action(observation, task_type)270 271 return self._parse_response(response_text, task_type, observation)272 273 def _parse_response(274 self, response_text: str, task_type: str, observation275 ) -> SecOpsAction:276 redacted_text = None277 public_resources = None278 fixed_resources = None279 ghost_users = None280 disabled_users = None281 classification = None282 severity = None283 config_issues = None284 hardened_config = None285 286 if task_type == "pii_redaction":287 redacted_text = self._extract_redacted_text(response_text)288 if not redacted_text:289 redacted_text = response_text.strip()290 291 elif task_type == "public_access":292 json_data = self._extract_json(response_text)293 if json_data:294 public_resources = json_data.get("identified_public", [])295 if not public_resources:296 public_resources = []297 fixed_resources = public_resources298 299 elif task_type == "ghost_user":300 json_data = self._extract_json(response_text)301 if json_data:302 ghost_users = json_data.get("identified_ghosts", [])303 if not ghost_users:304 ghost_users = []305 disabled_users = ghost_users306 307 elif task_type == "log_analysis":308 json_data = self._extract_json(response_text)309 if json_data:310 classification = json_data.get("classification")311 severity = json_data.get("severity")312 313 elif task_type == "config_hardening":314 json_data = self._extract_json(response_text)315 if json_data:316 config_issues = json_data.get("config_issues", [])317 hardened_config = json_data.get("hardened_config")318 319 reasoning = ""320 reasoning_match = re.search(r'"reasoning":\s*"([^"]*)"', response_text)321 if reasoning_match:322 reasoning = reasoning_match.group(1)323 324 return SecOpsAction(325 task_type=task_type,326 action_type=ActionType.FINALIZE,327 redacted_text=redacted_text,328 public_resources=public_resources,329 fixed_resources=fixed_resources,330 ghost_users=ghost_users,331 disabled_users=disabled_users,332 classification=classification,333 severity=severity,334 config_issues=config_issues,335 hardened_config=hardened_config,336 reasoning=reasoning,337 )338 339 def _fallback_action(self, observation, task_type: str) -> SecOpsAction:340 if task_type == "pii_redaction":341 text = observation.context.get("text", "")342 redacted = text343 expected_pii = observation.context.get("expected_pii", [])344 for pii in expected_pii:345 redacted = redacted.replace(pii.get("value", ""), "[REDACTED]")346 return SecOpsAction(347 task_type=task_type,348 action_type=ActionType.FINALIZE,349 redacted_text=redacted,350 )351 elif task_type == "public_access":352 public = [353 r["name"]354 for r in observation.context.get("resources", [])355 if r.get("public")356 ]357 return SecOpsAction(358 task_type=task_type,359 action_type=ActionType.FINALIZE,360 public_resources=public,361 fixed_resources=public,362 )363 elif task_type == "ghost_user":364 ghosts = [365 u["username"]366 for u in observation.context.get("users", [])367 if u.get("is_ghost")368 ]369 return SecOpsAction(370 task_type=task_type,371 action_type=ActionType.FINALIZE,372 ghost_users=ghosts,373 disabled_users=ghosts,374 )375 elif task_type == "log_analysis":376 return SecOpsAction(377 task_type=task_type,378 action_type=ActionType.FINALIZE,379 classification="NEEDS_INVESTIGATION",380 severity="MEDIUM",381 reasoning="Fallback: unable to analyze logs",382 )383 elif task_type == "config_hardening":384 return SecOpsAction(385 task_type=task_type,386 action_type=ActionType.FINALIZE,387 config_issues=[],388 hardened_config=observation.context.get("config_content", ""),389 reasoning="Fallback: unable to analyze config",390 )391 392 return SecOpsAction(task_type=task_type, action_type=ActionType.NOOP)393 394 395def run_episode(396 env: SecOpsEnv, agent: SecOpsAgent, task: str, difficulty: str397) -> tuple:398 rewards: List[float] = []399 steps_taken = 0400 error_msg = None401 402 try:403 result = env.reset(task=task, difficulty=difficulty)404 observation = result.observation405 406 for step in range(1, MAX_STEPS + 1):407 if result.done:408 break409 410 action = agent.generate_action(observation)411 action_str = str(action.model_dump())[:100]412 413 result = env.step(action)414 observation = result.observation415 416 reward = result.reward or 0.0417 done = result.done or False418 419 rewards.append(reward)420 steps_taken = step421 422 log_step(423 step=step, action=action_str, reward=reward, done=done, error=error_msg424 )425 426 if done:427 break428 429 score = observation.reward if observation.reward else sum(rewards)430 success = score >= SUCCESS_SCORE_THRESHOLD431 432 return success, steps_taken, score, rewards, error_msg433 434 except Exception as e:435 error_msg = str(e)[:100]436 if DEBUG:437 import traceback438 439 traceback.print_exc()440 score = sum(rewards) if rewards else 0.01441 return False, steps_taken, score, rewards, error_msg442 443 444def main():445 print(f"SecOps Environment Inference - Benchmark: {BENCHMARK}", flush=True)446 print(f"Model: {MODEL_NAME}", flush=True)447 print(f"API Base: {API_BASE_URL}", flush=True)448 449 if not HF_TOKEN:450 print("NOTE: HF_TOKEN not set - using fallback actions", flush=True)451 452 try:453 env = SecOpsEnv(base_url="http://localhost:8000")454 455 agent_config = AgentConfig(456 model_name=MODEL_NAME,457 api_base_url=API_BASE_URL,458 api_key=HF_TOKEN,459 )460 agent = SecOpsAgent(agent_config)461 462 for task_name, difficulty in ALL_TASKS:463 log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)464 465 success, steps, score, rewards, error = run_episode(466 env, agent, task_name, difficulty467 )468 469 log_end(success=success, steps=steps, score=score, rewards=rewards)470 471 env.close()472 473 except KeyboardInterrupt:474 print("\nBenchmark interrupted by user", flush=True)475 except Exception as e:476 print(f"\nError during benchmark: {e}", flush=True)477 if DEBUG:478 import traceback479 480 traceback.print_exc()481 482 483if __name__ == "__main__":484 main()485