""" server/environment.py — xsecure incident response simulation. Extends openenv-core Environment base class correctly. """ from __future__ import annotations import uuid from typing import Any, Dict, List, Optional try: from openenv_core.env_server import Environment except ImportError: from core.env_server import Environment from models import ( ActionType, IncidentAction, IncidentObservation, IncidentState, ) # --------------------------------------------------------------------------- # Task definitions # --------------------------------------------------------------------------- TASKS: Dict[int, dict] = { 1: { "attack_type": "brute_force", "attacker_ip": "192.168.1.47", "target_user": "admin", "target_service": "auth-service", "max_steps": 12, "logs": [ {"log_id": "L001", "message": "Multiple failed login attempts from 192.168.1.47 (47 attempts in 60s)", "timestamp": "2024-01-15T08:01:00Z"}, {"log_id": "L002", "message": "Account lockout triggered for user admin after failed logins from 192.168.1.47", "timestamp": "2024-01-15T08:01:45Z"}, {"log_id": "L003", "message": "Successful login from 192.168.1.47 after lockout bypass", "timestamp": "2024-01-15T08:03:10Z"}, ], "alerts": [ {"alert_id": "A001", "message": "Brute force pattern detected — 47 failed logins in 60 seconds", "severity": "high"}, {"alert_id": "A002", "message": "Account lockout bypass attempt detected", "severity": "high"}, ], "services": [ {"name": "auth-service", "status": "degraded"}, {"name": "web-app", "status": "running"}, {"name": "database", "status": "running"}, ], "active_users": ["admin", "alice", "bob"], "correct_sequence": { "analyze_log:L001": 0.15, "analyze_log:L002": 0.15, "block_ip:192.168.1.47": 0.50, }, "wrong_penalty": -0.15, "delay_penalty": 0.05, "speed_bonus_threshold": 6, "speed_bonus": 0.20, "compromise_steps": 8, }, 2: { "attack_type": "suspicious_login", "attacker_ip": "203.0.113.55", "target_user": "carol", "target_service": "hr-portal", "max_steps": 14, "logs": [ {"log_id": "L001", "message": "Login for carol from unusual geo-location (203.0.113.55 — Eastern Europe)", "timestamp": "2024-01-15T14:22:00Z"}, {"log_id": "L002", "message": "carol accessed sensitive HR records 3 minutes after login", "timestamp": "2024-01-15T14:25:10Z"}, {"log_id": "L003", "message": "carol attempted to export 1,200 employee records", "timestamp": "2024-01-15T14:27:33Z"}, ], "alerts": [ {"alert_id": "A001", "message": "Login from unusual location for carol", "severity": "medium"}, {"alert_id": "A002", "message": "Unusual data access pattern — bulk HR record access", "severity": "high"}, ], "services": [ {"name": "hr-portal", "status": "running"}, {"name": "auth-service", "status": "running"}, {"name": "database", "status": "running"}, ], "active_users": ["carol", "dave", "alice"], "correct_sequence": { "analyze_log:L001": 0.10, "analyze_log:L002": 0.10, "trace_user:carol": 0.20, "disable_account:carol": 0.45, }, "wrong_penalty": -0.15, "delay_penalty": 0.05, "speed_bonus_threshold": 7, "speed_bonus": 0.15, "compromise_steps": 10, }, 3: { "attack_type": "multi_stage", "attacker_ip": "198.51.100.23", "target_user": "eve", "target_service": "database", "max_steps": 18, "logs": [ {"log_id": "L001", "message": "Phishing email link clicked by eve — redirect to 198.51.100.23", "timestamp": "2024-01-15T09:00:00Z"}, {"log_id": "L002", "message": "Credential theft tool executed on eve's workstation", "timestamp": "2024-01-15T09:15:22Z"}, {"log_id": "L003", "message": "eve authenticated to database server outside normal hours", "timestamp": "2024-01-15T09:22:45Z"}, {"log_id": "L004", "message": "Lateral movement: eve's credentials used on prod-01 and database", "timestamp": "2024-01-15T09:30:11Z"}, {"log_id": "L005", "message": "Ransomware staging detected on database — encryption not yet started", "timestamp": "2024-01-15T09:45:00Z"}, ], "alerts": [ {"alert_id": "A001", "message": "Phishing link accessed — possible credential compromise for eve", "severity": "medium"}, {"alert_id": "A002", "message": "Credential harvesting tool detected", "severity": "high"}, {"alert_id": "A003", "message": "Lateral movement across prod-01 and database", "severity": "high"}, ], "services": [ {"name": "database", "status": "degraded"}, {"name": "prod-01", "status": "degraded"}, {"name": "web-app", "status": "running"}, ], "active_users": ["eve", "frank", "grace"], "correct_sequence": { "analyze_log:L001": 0.08, "analyze_log:L002": 0.08, "analyze_log:L004": 0.08, "trace_user:eve": 0.12, "block_ip:198.51.100.23": 0.20, "disable_account:eve": 0.20, "restart_service:database": 0.12, "restart_service:prod-01": 0.12, }, "wrong_penalty": -0.10, "delay_penalty": 0.04, "speed_bonus_threshold": 10, "speed_bonus": 0.10, "compromise_steps": 12, }, } # --------------------------------------------------------------------------- # Environment # --------------------------------------------------------------------------- class IncidentEnvironment(Environment): SUPPORTS_CONCURRENT_SESSIONS = True def __init__(self): super().__init__() self._state: Optional[IncidentState] = None self._task_def: Optional[dict] = None # ---- openenv-core interface ----------------------------------------- def reset( self, seed: Optional[int] = None, episode_id: Optional[str] = None, **kwargs: Any, ) -> IncidentObservation: task_id = int(kwargs.get("task_id", 1)) if task_id not in TASKS: task_id = 1 td = TASKS[task_id] self._task_def = td self._state = IncidentState( episode_id=episode_id or str(uuid.uuid4()), task_id=task_id, step_count=0, max_steps=td["max_steps"], attack_type=td["attack_type"], attacker_ip=td["attacker_ip"], target_user=td["target_user"], target_service=td["target_service"], ) return self._build_observation("Episode started. Investigate the alerts and logs.") def step( self, action: IncidentAction, timeout_s: Optional[float] = None, **kwargs: Any, ) -> IncidentObservation: if self._state is None: raise RuntimeError("Call reset() before step().") state = self._state td = self._task_def if state.done: return self._build_observation("Episode already finished.", reward=0.0) state.step_count += 1 action_key = f"{action.action_type}:{action.target}" reward = 0.0 result_msg = "" # Delay penalty every step reward -= td["delay_penalty"] atype = action.action_type if atype == ActionType.ANALYZE_LOG: r, result_msg = self._handle_analyze_log(action.target, action_key) reward += r elif atype == ActionType.TRACE_USER: r, result_msg = self._handle_trace_user(action.target, action_key) reward += r elif atype == ActionType.BLOCK_IP: r, result_msg = self._handle_block_ip(action.target, action_key) reward += r elif atype == ActionType.DISABLE_ACCOUNT: r, result_msg = self._handle_disable_account(action.target, action_key) reward += r elif atype == ActionType.RESTART_SERVICE: r, result_msg = self._handle_restart_service(action.target, action_key) reward += r elif atype == ActionType.IGNORE: state.progress_level = min(state.progress_level + 1, state.max_progress) reward -= 0.10 result_msg = "No action taken. Attack progresses!" else: state.wrong_actions += 1 reward += td["wrong_penalty"] result_msg = f"Unknown action: {atype}" reward = round(reward, 4) state.total_reward += reward done, info = self._check_termination() if done: state.done = True return self._build_observation(result_msg, reward=reward, done=done, info=info) @property def state(self) -> IncidentState: if self._state is None: raise RuntimeError("Call reset() first.") return self._state # ---- Action handlers ----------------------------------------------- def _handle_analyze_log(self, target: str, key: str): state = self._state td = self._task_def log_ids = [l["log_id"] for l in td["logs"]] if target not in log_ids: state.wrong_actions += 1 return td["wrong_penalty"], f"Log {target} does not exist." if target in state.revealed_logs: return -0.05, f"Log {target} already analyzed." state.revealed_logs.append(target) if key in td["correct_sequence"]: state.correct_detections += 1 msg = next(l["message"] for l in td["logs"] if l["log_id"] == target) return td["correct_sequence"][key], f"[ANALYSIS] {target}: {msg}" return 0.05, f"Log {target} analyzed — no significant findings." def _handle_trace_user(self, target: str, key: str): state = self._state td = self._task_def if target not in td["active_users"]: state.wrong_actions += 1 return td["wrong_penalty"], f"User {target!r} not found." if target in state.revealed_users: return -0.05, f"User {target} already traced." state.revealed_users.append(target) if key in td["correct_sequence"]: state.correct_detections += 1 return td["correct_sequence"][key], f"[TRACE] {target}: Confirmed suspicious activity." return 0.05, f"User {target} traced — activity appears normal." def _handle_block_ip(self, target: str, key: str): state = self._state td = self._task_def if target in state.blocked_ips: return -0.05, f"IP {target} already blocked." state.blocked_ips.append(target) if key in td["correct_sequence"]: state.correct_detections += 1 multiplier = 1.0 if state.correct_detections > 1 else 0.6 return td["correct_sequence"][key] * multiplier, f"[BLOCKED] IP {target} blocked." state.wrong_actions += 1 return td["wrong_penalty"], f"Blocking {target} was incorrect." def _handle_disable_account(self, target: str, key: str): state = self._state td = self._task_def if target in state.disabled_accounts: return -0.05, f"Account {target} already disabled." state.disabled_accounts.append(target) if key in td["correct_sequence"]: state.correct_detections += 1 multiplier = 1.0 if state.correct_detections > 1 else 0.6 return td["correct_sequence"][key] * multiplier, f"[DISABLED] Account {target} disabled." state.wrong_actions += 1 return td["wrong_penalty"], f"Disabling {target} was incorrect." def _handle_restart_service(self, target: str, key: str): state = self._state td = self._task_def valid = [s["name"] for s in td["services"]] if target not in valid: state.wrong_actions += 1 return td["wrong_penalty"], f"Service {target!r} not found." if target in state.restarted_services: return -0.05, f"Service {target} already restarted." state.restarted_services.append(target) if key in td["correct_sequence"]: state.correct_detections += 1 return td["correct_sequence"][key], f"[RESTARTED] Service {target} restored." state.wrong_actions += 1 return td["wrong_penalty"], f"Restarting {target} was not necessary." # ---- Termination --------------------------------------------------- def _check_termination(self): state = self._state td = self._task_def seq = td["correct_sequence"] all_done = all(self._action_completed(k) for k in seq) speed_bonus = 0.0 if all_done and state.step_count <= td["speed_bonus_threshold"]: speed_bonus = td["speed_bonus"] state.total_reward += speed_bonus compromise = ( state.progress_level >= state.max_progress or state.step_count >= td["compromise_steps"] ) and not all_done timeout = state.step_count >= state.max_steps done = all_done or compromise or timeout if done: state.success = all_done and not compromise state.compromise = compromise max_possible = sum(seq.values()) + td["speed_bonus"] final_score = round(min(max(state.total_reward, 0.0) / max_possible, 1.0), 4) if max_possible > 0 else 0.0 info = { "final_score": final_score, "success": state.success, "compromise": state.compromise, "wrong_actions": state.wrong_actions, "speed_bonus": speed_bonus, } if done else {} return done, info def _action_completed(self, key: str) -> bool: action_type, target = key.split(":", 1) s = self._state return { "analyze_log": target in s.revealed_logs, "trace_user": target in s.revealed_users, "block_ip": target in s.blocked_ips, "disable_account": target in s.disabled_accounts, "restart_service": target in s.restarted_services, }.get(action_type, False) # ---- Observation builder ------------------------------------------- def _build_observation( self, result_msg: str, reward: float = 0.0, done: bool = False, info: dict = None, ) -> IncidentObservation: td = self._task_def return IncidentObservation( logs=td["logs"], alerts=td["alerts"], services=td["services"], active_users=td["active_users"], step_count=self._state.step_count, reward=reward, done=done, info=info or {}, last_action_result=result_msg, )