| |
| import random |
| from typing import Any, Dict, List, Optional |
|
|
| from models import ( |
| Action, |
| ActionType, |
| LogEntry, |
| Observation, |
| ObservationMetadata, |
| Reward, |
| RewardDetails, |
| StepInfo, |
| ) |
| from tasks import BaseTask, BruteForceSSHTask, LateralMovementTask, APTMultiStageTask |
|
|
| |
| |
| |
| class AttackState: |
| """ |
| Hidden truth the agent never sees. |
| The fields below are deliberately mutable only inside the env. |
| """ |
|
|
| stage: str |
| compromised_hosts: List[str] |
| attacker_ips: List[str] |
| data_exfiltrated: bool |
|
|
| |
| pivot_host: Optional[str] = None |
| internal_host: Optional[str] = None |
|
|
| def __init__(self) -> None: |
| self.stage = "initial_access" |
| self.compromised_hosts = [] |
| self.attacker_ips = [] |
| self.data_exfiltrated = False |
| self.pivot_host = None |
| self.internal_host = None |
|
|
| |
| def advance_stage(self) -> None: |
| if self.stage == "initial_access": |
| self.stage = "lateral_movement" |
| elif self.stage == "lateral_movement": |
| self.stage = "exfiltration" |
|
|
| def is_finished(self) -> bool: |
| return self.data_exfiltrated |
|
|
| |
| |
| |
| class AIGymEnv: |
| MAX_STEPS = 20 |
| NOISE_SOURCES = ["ssh", "http", "dns"] |
|
|
| def __init__(self, seed: int | None = None) -> None: |
| self._rand = random.Random(seed) |
| self._step_counter: int = 0 |
| self._state: AttackState = AttackState() |
| self._blocked_ips: set[str] = set() |
| self._isolated_hosts: set[str] = set() |
| self._last_action_effect: Optional[str] = None |
| self._task: Optional[BaseTask] = None |
| self._ip_pool = [ |
| f"10.0.{i}.{j}" for i in range(1, 5) for j in range(1, 255) |
| ] |
|
|
| |
| |
| |
| def reset(self) -> Observation: |
| self._step_counter = 0 |
| self._state = AttackState() |
| self._blocked_ips.clear() |
| self._isolated_hosts.clear() |
| if self._task: |
| self._task.initialize(self) |
| init_logs = self._generate_logs(phase="reset") + self._task_advance() |
| return Observation( |
| logs=init_logs, |
| metadata=ObservationMetadata(step=self._step_counter, alerts_triggered=0), |
| ) |
|
|
| def step(self, action: Action) -> tuple[Observation, Reward, bool, StepInfo]: |
| self._step_counter += 1 |
| valid, msg = self._validate_action(action) |
| if not valid: |
| reward = self._compute_reward(detection=0.05, false_positive=0.95, efficiency=0.05) |
| info = StepInfo(reason=f"Invalid action: {msg}", confidence=0.05, action_effect="none") |
| obs = Observation( |
| logs=self._generate_logs() + self._task_advance(), |
| metadata=ObservationMetadata(step=self._step_counter, alerts_triggered=0), |
| ) |
| return obs, reward, False, info |
|
|
| self._apply_action(action) |
| mitigated = self._task_mitigated() |
| if not mitigated: |
| self._state.advance_stage() |
| logs = self._generate_logs() + self._task_advance() |
| observation = Observation( |
| logs=logs, |
| metadata=ObservationMetadata( |
| step=self._step_counter, |
| alerts_triggered=self._count_alerts(logs), |
| ), |
| ) |
| done = ( |
| self._step_counter >= self.MAX_STEPS or self._state.is_finished() or mitigated |
| ) |
| |
| if done: |
| detection = 0.95 if mitigated else 0.05 |
| false_positive = 0.95 if self._action_was_false_positive(action) else 0.05 |
| efficiency = max(0.05, min(0.95, 1.0 - (self._step_counter / self.MAX_STEPS))) |
| reward = self._compute_reward(detection, false_positive, efficiency, apply_bonus=True) |
| else: |
| reward = self._compute_reward(0.05, 0.95, 0.05, apply_bonus=False) |
| info = StepInfo( |
| reason=self._explain_reason(), |
| confidence=self._rand.uniform(0.7, 0.95), |
| action_effect=self._last_action_effect, |
| ) |
| return observation, reward, done, info |
|
|
| def state(self) -> Dict[str, Any]: |
| """Return a serializable snapshot of the current environment state. |
| |
| Required by the OpenEnv specification. |
| """ |
| return { |
| "step": self._step_counter, |
| "stage": self._state.stage, |
| "attacker_ips": list(self._state.attacker_ips), |
| "compromised_hosts": list(self._state.compromised_hosts), |
| "pivot_host": self._state.pivot_host, |
| "internal_host": self._state.internal_host, |
| "data_exfiltrated": self._state.data_exfiltrated, |
| "blocked_ips": list(self._blocked_ips), |
| "isolated_hosts": list(self._isolated_hosts), |
| } |
|
|
| |
| |
| |
| def load_task(self, task: BaseTask) -> None: |
| self._task = task |
| task.initialize(self) |
|
|
| def _task_advance(self) -> List[LogEntry]: |
| if self._task: |
| return self._task.advance(self) |
| return [] |
|
|
| def _task_mitigated(self) -> bool: |
| if self._task and hasattr(self._task, "is_mitigated"): |
| return self._task.is_mitigated(self) |
| return False |
|
|
| |
| |
| |
| def _validate_action(self, action: Action) -> tuple[bool, str]: |
| if action.type == ActionType.BLOCK_IP: |
| if action.target_type != "ip": |
| return False, "BLOCK_IP requires target_type='ip'" |
| if action.target not in self._ip_pool and action.target not in self._state.attacker_ips: |
| return False, "IP unknown to the system" |
| if action.target in self._blocked_ips: |
| return False, "IP already blocked" |
| elif action.type == ActionType.ISOLATE_HOST: |
| if action.target_type != "host": |
| return False, "ISOLATE_HOST requires target_type='host'" |
| elif action.type == ActionType.ALLOW: |
| if action.target_type not in {"ip", "host"}: |
| return False, "ALLOW must target ip or host" |
| return True, "" |
|
|
| def _apply_action(self, action: Action) -> None: |
| if action.type == ActionType.BLOCK_IP: |
| self._blocked_ips.add(action.target) |
| self._last_action_effect = f"IP {action.target} blocked" |
| elif action.type == ActionType.ISOLATE_HOST: |
| self._isolated_hosts.add(action.target) |
| self._last_action_effect = f"Host {action.target} isolated" |
| elif action.type == ActionType.ALLOW: |
| self._last_action_effect = ( |
| f"Allowed {action.target_type} {action.target}" |
| ) |
| elif action.type == ActionType.INVESTIGATE: |
| if ( |
| isinstance(getattr(self, "_task", None), BruteForceSSHTask) |
| and action.target_type == "ip" |
| and action.target == self._state.attacker_ips[0] |
| ): |
| self._task.investigation_done = True |
| if ( |
| isinstance(getattr(self, "_task", None), LateralMovementTask) |
| and action.target_type == "host" |
| and action.target in {self._state.pivot_host, self._state.internal_host} |
| ): |
| self._task.hosts_investigated.add(action.target) |
| if ( |
| isinstance(getattr(self, "_task", None), APTMultiStageTask) |
| and action.target_type in {"host", "ip"} |
| ): |
| self._task.hosts_investigated.add(action.target) |
| self._last_action_effect = f"Investigated {action.target_type} {action.target}" |
| else: |
| self._last_action_effect = "no effect" |
|
|
| |
| def _action_was_false_positive(self, action: Action) -> bool: |
| if action.type == ActionType.BLOCK_IP: |
| if action.target not in self._state.attacker_ips: |
| return True |
| if action.type == ActionType.ISOLATE_HOST: |
| if ( |
| action.target not in self._state.compromised_hosts |
| and action.target != self._state.pivot_host |
| and action.target != self._state.internal_host |
| ): |
| return True |
| return False |
|
|
| |
| def _generate_logs(self, phase: str = "step") -> List[LogEntry]: |
| logs: List[LogEntry] = [] |
| for _ in range(self._rand.randint(2, 5)): |
| src = self._rand.choice(self.NOISE_SOURCES) |
| logs.append( |
| LogEntry( |
| timestamp=self._rand_timestamp(), |
| source=src, |
| severity=self._rand.choice(["INFO", "WARN"]), |
| message=self._random_noise_message(src), |
| ip=self._rand.choice(self._ip_pool) if src != "dns" else None, |
| user=self._rand.choice(["alice", "bob", "carol"]), |
| hostname=f"host-{self._rand.randint(1,20)}.corp", |
| event_type=self._rand.choice( |
| ["login", "file_access", "network_conn"] |
| ), |
| ) |
| ) |
| return logs |
|
|
| def _rand_timestamp(self) -> str: |
| base = 1672502400 |
| ts = base + self._step_counter * 60 + self._rand.randint(0, 30) |
| return f"2026-04-07T{(ts % 86400)//3600:02d}:{(ts % 3600)//60:02d}:{ts % 60:02d}Z" |
|
|
| def _random_noise_message(self, src: str) -> str: |
| if src == "ssh": |
| return ( |
| f"Accepted publickey for {self._rand.choice(['alice', 'bob'])} " |
| f"from {self._rand.choice(self._ip_pool)}" |
| ) |
| if src == "http": |
| return "GET /api/v1/resource HTTP/1.1 200 OK" |
| if src == "dns": |
| return "Query A example.com" |
| return "Generic log entry" |
|
|
| def _count_alerts(self, logs: List[LogEntry]) -> int: |
| return sum(1 for l in logs if l.severity in {"WARN", "ERROR", "CRITICAL"}) |
|
|
| |
| def _compute_reward( |
| self, |
| detection: float, |
| false_positive: float, |
| efficiency: float, |
| apply_bonus: bool = True, |
| ) -> Reward: |
| from graders import compute_reward |
| bonus = 0.0 |
| if apply_bonus: |
| if isinstance(getattr(self, "_task", None), LateralMovementTask): |
| task: LateralMovementTask = self._task |
| if task.hosts_investigated: |
| bonus = 0.04 |
| if isinstance(getattr(self, "_task", None), APTMultiStageTask): |
| apt_task: APTMultiStageTask = self._task |
| if apt_task.hosts_investigated: |
| bonus = 0.06 |
| return compute_reward(detection, false_positive, efficiency, investigation_bonus=bonus) |
|
|
| |
| def _explain_reason(self) -> str: |
| if self._state.stage == "initial_access": |
| return "Brute-force attempts detected" |
| if self._state.stage == "lateral_movement": |
| return "Suspicious internal traffic observed" |
| if self._state.stage == "exfiltration": |
| return "Large outbound data transfer observed" |
| return "Noise / normal operation" |
|
|
| |
| def dump_internal_state(self) -> str: |
| return ( |
| f"Step: {self._step_counter}\n" |
| f"Stage: {self._state.stage}\n" |
| f"Attacker IPs: {self._state.attacker_ips}\n" |
| f"Compromised hosts: {self._state.compromised_hosts}\n" |
| f"Pivot host: {self._state.pivot_host}\n" |
| f"Internal host: {self._state.internal_host}\n" |
| f"Blocked IPs: {list(self._blocked_ips)}\n" |
| f"Isolated hosts: {list(self._isolated_hosts)}\n" |
| f"Data exfiltrated: {self._state.data_exfiltrated}" |
| ) |
|
|
