Spaces:

netha01
/

SOCEnv

Sleeping

SOCEnv / tasks.py

Add app.py and YAML configuration

6edd399 7 months ago

20.7 kB

# tasks.py – scenario definitions for the AI SOC Gym
"""
This module defines the attack scenarios (tasks) that drive the hidden
`AttackState` inside `environment.py`.  Each task implements a tiny
interface (`BaseTask`) used by the environment to generate logs and to
report whether the analyst has successfully mitigated the attack.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, List, Set

from models import Action, ActionType, LogEntry

if TYPE_CHECKING:
    from environment import AIGymEnv

# ----------------------------------------------------------------------
# BaseTask – minimal contract for every scenario
# ----------------------------------------------------------------------
class BaseTask:
    def initialize(self, env: AIGymEnv) -> None:
        """Set up hidden state before the first reset."""
        raise NotImplementedError

    def advance(self, env: AIGymEnv) -> List[LogEntry]:
        """Return task‑specific logs for the current step."""
        raise NotImplementedError

    # Most tasks will override this; default = never mitigated.
    def is_mitigated(self, env: AIGymEnv) -> bool:  # pragma: no cover
        return False

# ----------------------------------------------------------------------
# Task 1 – Brute‑Force SSH (easy)
# ----------------------------------------------------------------------
@dataclass
class BruteForceSSHTask(BaseTask):
    attacker_ip: str = "203.0.113.5"
    max_failures: int = 4
    investigation_done: bool = False
    _failed: int = field(default=0, init=False)
    _success: bool = field(default=False, init=False)

    def initialize(self, env: AIGymEnv) -> None:
        env._state.stage = "initial_access"
        env._state.attacker_ips = [self.attacker_ip]
        self._failed = 0
        self._success = False
        self.investigation_done = False

    def advance(self, env: AIGymEnv) -> List[LogEntry]:
        logs: List[LogEntry] = []
        if not self._success:
            if self._failed < self.max_failures:
                logs.append(
                    LogEntry(
                        timestamp=env._rand_timestamp(),
                        source="ssh",
                        severity="WARN",
                        message=f"Failed password for invalid user admin from {self.attacker_ip}",
                        ip=self.attacker_ip,
                        user="admin",
                        hostname="gateway.corp",
                        event_type="login",
                    )
                )
                self._failed += 1
                return logs
            # success path – attacker gains foothold on gateway
            logs.append(
                LogEntry(
                    timestamp=env._rand_timestamp(),
                    source="ssh",
                    severity="INFO",
                    message=f"Accepted password for root from {self.attacker_ip}",
                    ip=self.attacker_ip,
                    user="root",
                    hostname="gateway.corp",
                    event_type="login",
                )
            )
            self._success = True
            env._state.compromised_hosts.append("gateway.corp")
            env._state.data_exfiltrated = True
        return logs

    def is_mitigated(self, env: AIGymEnv) -> bool:
        return self.attacker_ip in env._blocked_ips

# ----------------------------------------------------------------------
# Task 2 – Lateral Movement (medium difficulty)
# ----------------------------------------------------------------------
@dataclass
class LateralMovementTask(BaseTask):
    """Medium scenario – external brute‑force, pivot host, internal hop.
    Steps:
    1️⃣ Brute‑force external login.
    2️⃣ Compromise **pivot_host** and SSH to **internal_host**.
    3️⃣ Create admin on **internal_host**.
    The analyst must eventually isolate a compromised host.
    """

    attacker_ip: str = "203.0.113.5"
    max_failures: int = 3
    pivot_host: str = "host-12.corp"
    internal_host: str = "host-27.corp"

    # internal counters – kept on the task instance
    _external_failed: int = field(default=0, init=False)
    _internal_failed: int = field(default=0, init=False)
    _pivot_done: bool = field(default=False, init=False)
    _internal_done: bool = field(default=False, init=False)
    hosts_investigated: Set[str] = field(default_factory=set, init=False)

    def initialize(self, env: AIGymEnv) -> None:
        env._state.stage = "initial_access"
        env._state.attacker_ips = [self.attacker_ip]
        env._state.pivot_host = self.pivot_host
        env._state.internal_host = self.internal_host
        self._external_failed = 0
        self._internal_failed = 0
        self._pivot_done = False
        self._internal_done = False
        self.hosts_investigated = set()

    def advance(self, env: AIGymEnv) -> List[LogEntry]:
        logs: List[LogEntry] = []
        # ----- Phase 1 – external brute‑force -----
        if not self._pivot_done:
            if self._external_failed < self.max_failures:
                logs.append(
                    LogEntry(
                        timestamp=env._rand_timestamp(),
                        source="ssh",
                        severity="WARN",
                        message=f"Failed password for invalid user admin from {self.attacker_ip}",
                        ip=self.attacker_ip,
                        user="admin",
                        hostname="gateway.corp",
                        event_type="login",
                    )
                )
                self._external_failed += 1
                return logs
            # success – pivot host compromised
            logs.append(
                LogEntry(
                    timestamp=env._rand_timestamp(),
                    source="ssh",
                    severity="INFO",
                    message=f"Accepted password for svc_user from {self.attacker_ip} on {self.pivot_host}",
                    ip=self.attacker_ip,
                    user="svc_user",
                    hostname=self.pivot_host,
                    event_type="login",
                )
            )
            env._state.compromised_hosts.append(self.pivot_host)
            self._pivot_done = True
            return logs

        # ----- Phase 2 – internal SSH hop -----
        if self._pivot_done and not self._internal_done:
            logs.append(
                LogEntry(
                    timestamp=env._rand_timestamp(),
                    source="ssh",
                    severity="INFO",
                    message=f"{self.pivot_host} initiated SSH connection to {self.internal_host}",
                    ip=self.pivot_host,
                    user="svc_user",
                    hostname=self.internal_host,
                    event_type="network_conn",
                )
            )
            # after a couple of steps the internal host gets admin rights
            if self._internal_failed >= 2:
                logs.append(
                    LogEntry(
                        timestamp=env._rand_timestamp(),
                        source="ssh",
                        severity="INFO",
                        message=f"New admin account created on {self.internal_host} by svc_user",
                        ip=self.internal_host,
                        user="admin",
                        hostname=self.internal_host,
                        event_type="login",
                    )
                )
                env._state.compromised_hosts.append(self.internal_host)
                self._internal_done = True
                env._state.data_exfiltrated = True
            else:
                self._internal_failed += 1
            return logs

        # Phase 3 – attack already succeeded – nothing extra to emit
        return logs

    def is_mitigated(self, env: AIGymEnv) -> bool:
        compromised = set(env._state.compromised_hosts)
        isolated = set(env._isolated_hosts)
        return len(compromised & isolated) > 0


# ----------------------------------------------------------------------
# Task 3 – APT Multi‑Stage (hard difficulty)
# ----------------------------------------------------------------------
@dataclass
class APTMultiStageTask(BaseTask):
    """Hard scenario — a full Advanced Persistent Threat kill-chain.
    Phases:
    1️⃣  **Phishing** — Spear‑phishing email delivers a malicious document
         to a user workstation.  The logs show an inbound email with a
         suspicious attachment, followed by an outbound C2 beacon.
    2️⃣  **Credential Theft** — The compromised workstation runs a
         credential‑dumping tool (Mimikatz‑style).  Logs show LSASS
         access and new logon events with harvested service accounts.
    3️⃣  **Lateral Movement** — Stolen creds are used to hop through
         two internal hosts.  Logs show WinRM / SSH connections between
         internal machines.
    4️⃣  **Exfiltration** — Attacker stages data on a staging host and
         tunnels it out via DNS over HTTPS.  Logs show large DNS‑TXT
         responses and unusual outbound traffic volume.
    Mitigation requires the analyst to:
      • Isolate the phished workstation  **AND**
      • Block the C2 IP
    Partial credit is given for blocking C2 alone or isolating any
    compromised host.
    """

    # ---------- configurable knobs ----------
    c2_ip: str = "198.51.100.23"
    phish_target_user: str = "jdoe"
    phished_host: str = "ws-PC042.corp"
    staging_host: str = "file-srv-01.corp"
    dc_host: str = "dc-01.corp"
    max_phish_steps: int = 2  # steps before the user clicks the link

    # ---------- internal counters ----------
    _phase: int = field(default=0, init=False)          # 0–4
    _phase_counter: int = field(default=0, init=False)   # steps within phase
    _user_clicked: bool = field(default=False, init=False)
    _creds_dumped: bool = field(default=False, init=False)
    _lateral_hops: int = field(default=0, init=False)
    _exfil_started: bool = field(default=False, init=False)
    hosts_investigated: Set[str] = field(default_factory=set, init=False)

    def initialize(self, env: AIGymEnv) -> None:
        env._state.stage = "initial_access"
        env._state.attacker_ips = [self.c2_ip]
        env._state.pivot_host = self.phished_host
        env._state.internal_host = self.staging_host
        self._phase = 0
        self._phase_counter = 0
        self._user_clicked = False
        self._creds_dumped = False
        self._lateral_hops = 0
        self._exfil_started = False
        self.hosts_investigated = set()

    # ------------------------------------------------------------------
    def advance(self, env: AIGymEnv) -> List[LogEntry]:
        logs: List[LogEntry] = []
        self._phase_counter += 1

        # ---- Phase 0 : Phishing delivery ----
        if self._phase == 0:
            logs.append(
                LogEntry(
                    timestamp=env._rand_timestamp(),
                    source="email",
                    severity="INFO",
                    message=(
                        f"Inbound email to {self.phish_target_user}@corp.local "
                        f"from hr-benefits@secure‑docs.com "
                        f"subject='Q1 Compensation Update' attachment=Q1_Update.xlsm"
                    ),
                    ip=self.c2_ip,
                    user=self.phish_target_user,
                    hostname=self.phished_host,
                    event_type="email",
                )
            )
            if self._phase_counter >= self.max_phish_steps:
                # User clicks the doc — macro fires
                logs.append(
                    LogEntry(
                        timestamp=env._rand_timestamp(),
                        source="endpoint",
                        severity="WARN",
                        message=(
                            f"EXCEL.EXE spawned powershell.exe on {self.phished_host} "
                            f"(user={self.phish_target_user})"
                        ),
                        ip=None,
                        user=self.phish_target_user,
                        hostname=self.phished_host,
                        event_type="process",
                    )
                )
                logs.append(
                    LogEntry(
                        timestamp=env._rand_timestamp(),
                        source="firewall",
                        severity="WARN",
                        message=(
                            f"Outbound HTTPS connection from {self.phished_host} "
                            f"to {self.c2_ip}:443 (TLS SNI: cdn‑assets.com)"
                        ),
                        ip=self.c2_ip,
                        user=self.phish_target_user,
                        hostname=self.phished_host,
                        event_type="network_conn",
                    )
                )
                self._user_clicked = True
                env._state.compromised_hosts.append(self.phished_host)
                self._phase = 1
                self._phase_counter = 0
            return logs

        # ---- Phase 1 : Credential theft ----
        if self._phase == 1:
            # First step in this phase — LSASS access
            if not self._creds_dumped:
                logs.append(
                    LogEntry(
                        timestamp=env._rand_timestamp(),
                        source="endpoint",
                        severity="CRITICAL",
                        message=(
                            f"Suspicious LSASS memory access on {self.phished_host} "
                            f"by process rundll32.exe (PID 4728)"
                        ),
                        ip=None,
                        user="SYSTEM",
                        hostname=self.phished_host,
                        event_type="process",
                    )
                )
                logs.append(
                    LogEntry(
                        timestamp=env._rand_timestamp(),
                        source="auth",
                        severity="WARN",
                        message=(
                            f"New logon: svc_backup@corp.local on {self.phished_host} "
                            f"(logon type 9 — NewCredentials)"
                        ),
                        ip=None,
                        user="svc_backup",
                        hostname=self.phished_host,
                        event_type="login",
                    )
                )
                self._creds_dumped = True
                env._state.stage = "lateral_movement"
                self._phase = 2
                self._phase_counter = 0
            return logs

        # ---- Phase 2 : Lateral movement (2 hops) ----
        if self._phase == 2:
            if self._lateral_hops == 0:
                # Hop 1: phished_host → dc_host
                logs.append(
                    LogEntry(
                        timestamp=env._rand_timestamp(),
                        source="auth",
                        severity="WARN",
                        message=(
                            f"Kerberos TGS request for host/{self.dc_host} "
                            f"from {self.phished_host} (user=svc_backup)"
                        ),
                        ip=self.phished_host,
                        user="svc_backup",
                        hostname=self.dc_host,
                        event_type="login",
                    )
                )
                logs.append(
                    LogEntry(
                        timestamp=env._rand_timestamp(),
                        source="ssh",
                        severity="INFO",
                        message=(
                            f"WinRM session opened on {self.dc_host} "
                            f"by svc_backup from {self.phished_host}"
                        ),
                        ip=self.phished_host,
                        user="svc_backup",
                        hostname=self.dc_host,
                        event_type="network_conn",
                    )
                )
                env._state.compromised_hosts.append(self.dc_host)
                self._lateral_hops = 1
                return logs

            if self._lateral_hops == 1:
                # Hop 2: dc_host → staging_host
                logs.append(
                    LogEntry(
                        timestamp=env._rand_timestamp(),
                        source="ssh",
                        severity="INFO",
                        message=(
                            f"{self.dc_host} initiated SMB connection to "
                            f"{self.staging_host} (share: \\\\{self.staging_host}\\data$)"
                        ),
                        ip=self.dc_host,
                        user="svc_backup",
                        hostname=self.staging_host,
                        event_type="network_conn",
                    )
                )
                logs.append(
                    LogEntry(
                        timestamp=env._rand_timestamp(),
                        source="endpoint",
                        severity="WARN",
                        message=(
                            f"Batch file copy: 847 files (2.3 GB) staged in "
                            f"C:\\Windows\\Temp\\~cache on {self.staging_host}"
                        ),
                        ip=None,
                        user="svc_backup",
                        hostname=self.staging_host,
                        event_type="file_access",
                    )
                )
                env._state.compromised_hosts.append(self.staging_host)
                self._lateral_hops = 2
                env._state.stage = "exfiltration"
                self._phase = 3
                self._phase_counter = 0
                return logs

            return logs

        # ---- Phase 3 : Exfiltration via DNS tunneling ----
        if self._phase == 3:
            if not self._exfil_started:
                logs.append(
                    LogEntry(
                        timestamp=env._rand_timestamp(),
                        source="dns",
                        severity="WARN",
                        message=(
                            f"Anomalous DNS-over-HTTPS traffic from {self.staging_host} "
                            f"to {self.c2_ip} — 340 TXT queries in 60 s "
                            f"(avg payload 220 bytes)"
                        ),
                        ip=self.c2_ip,
                        user=None,
                        hostname=self.staging_host,
                        event_type="network_conn",
                    )
                )
                logs.append(
                    LogEntry(
                        timestamp=env._rand_timestamp(),
                        source="firewall",
                        severity="CRITICAL",
                        message=(
                            f"High-volume outbound: {self.staging_host} → "
                            f"{self.c2_ip}:443  |  2.1 GB transferred in 180 s"
                        ),
                        ip=self.c2_ip,
                        user=None,
                        hostname=self.staging_host,
                        event_type="network_conn",
                    )
                )
                self._exfil_started = True

            if self._phase_counter >= 2:
                env._state.data_exfiltrated = True

            return logs

        return logs

    # ------------------------------------------------------------------
    def is_mitigated(self, env: AIGymEnv) -> bool:
        """Full mitigation requires BOTH:
        1. The C2 IP is blocked  OR  the phished host is isolated.
        2. At least one compromised host is isolated.
        """
        c2_blocked = self.c2_ip in env._blocked_ips
        phished_isolated = self.phished_host in env._isolated_hosts
        any_compromised_isolated = bool(
            set(env._state.compromised_hosts) & set(env._isolated_hosts)
        )
        return (c2_blocked or phished_isolated) and any_compromised_isolated


# ----------------------------------------------------------------------
# Exported symbols for `import *`
# ----------------------------------------------------------------------
__all__ = [
    "BaseTask",
    "BruteForceSSHTask",
    "LateralMovementTask",
    "APTMultiStageTask",
]

	# tasks.py – scenario definitions for the AI SOC Gym
	"""
	This module defines the attack scenarios (tasks) that drive the hidden
	`AttackState` inside `environment.py`. Each task implements a tiny
	interface (`BaseTask`) used by the environment to generate logs and to
	report whether the analyst has successfully mitigated the attack.
	"""

	from __future__ import annotations

	from dataclasses import dataclass, field
	from typing import TYPE_CHECKING, List, Set

	from models import Action, ActionType, LogEntry

	if TYPE_CHECKING:
	from environment import AIGymEnv

	# ----------------------------------------------------------------------
	# BaseTask – minimal contract for every scenario
	# ----------------------------------------------------------------------
	class BaseTask:
	def initialize(self, env: AIGymEnv) -> None:
	"""Set up hidden state before the first reset."""
	raise NotImplementedError

	def advance(self, env: AIGymEnv) -> List[LogEntry]:
	"""Return task‑specific logs for the current step."""
	raise NotImplementedError

	# Most tasks will override this; default = never mitigated.
	def is_mitigated(self, env: AIGymEnv) -> bool: # pragma: no cover
	return False

	# ----------------------------------------------------------------------
	# Task 1 – Brute‑Force SSH (easy)
	# ----------------------------------------------------------------------
	@dataclass
	class BruteForceSSHTask(BaseTask):
	attacker_ip: str = "203.0.113.5"
	max_failures: int = 4
	investigation_done: bool = False
	_failed: int = field(default=0, init=False)
	_success: bool = field(default=False, init=False)

	def initialize(self, env: AIGymEnv) -> None:
	env._state.stage = "initial_access"
	env._state.attacker_ips = [self.attacker_ip]
	self._failed = 0
	self._success = False
	self.investigation_done = False

	def advance(self, env: AIGymEnv) -> List[LogEntry]:
	logs: List[LogEntry] = []
	if not self._success:
	if self._failed < self.max_failures:
	logs.append(
	LogEntry(
	timestamp=env._rand_timestamp(),
	source="ssh",
	severity="WARN",
	message=f"Failed password for invalid user admin from {self.attacker_ip}",
	ip=self.attacker_ip,
	user="admin",
	hostname="gateway.corp",
	event_type="login",
	)
	)
	self._failed += 1
	return logs
	# success path – attacker gains foothold on gateway
	logs.append(
	LogEntry(
	timestamp=env._rand_timestamp(),
	source="ssh",
	severity="INFO",
	message=f"Accepted password for root from {self.attacker_ip}",
	ip=self.attacker_ip,
	user="root",
	hostname="gateway.corp",
	event_type="login",
	)
	)
	self._success = True
	env._state.compromised_hosts.append("gateway.corp")
	env._state.data_exfiltrated = True
	return logs

	def is_mitigated(self, env: AIGymEnv) -> bool:
	return self.attacker_ip in env._blocked_ips

	# ----------------------------------------------------------------------
	# Task 2 – Lateral Movement (medium difficulty)
	# ----------------------------------------------------------------------
	@dataclass
	class LateralMovementTask(BaseTask):
	"""Medium scenario – external brute‑force, pivot host, internal hop.

	Steps:
	1️⃣ Brute‑force external login.
	2️⃣ Compromise pivot_host and SSH to internal_host.
	3️⃣ Create admin on internal_host.
	The analyst must eventually isolate a compromised host.
	"""

	attacker_ip: str = "203.0.113.5"
	max_failures: int = 3
	pivot_host: str = "host-12.corp"
	internal_host: str = "host-27.corp"

	# internal counters – kept on the task instance
	_external_failed: int = field(default=0, init=False)
	_internal_failed: int = field(default=0, init=False)
	_pivot_done: bool = field(default=False, init=False)
	_internal_done: bool = field(default=False, init=False)
	hosts_investigated: Set[str] = field(default_factory=set, init=False)

	def initialize(self, env: AIGymEnv) -> None:
	env._state.stage = "initial_access"
	env._state.attacker_ips = [self.attacker_ip]
	env._state.pivot_host = self.pivot_host
	env._state.internal_host = self.internal_host
	self._external_failed = 0
	self._internal_failed = 0
	self._pivot_done = False
	self._internal_done = False
	self.hosts_investigated = set()

	def advance(self, env: AIGymEnv) -> List[LogEntry]:
	logs: List[LogEntry] = []
	# ----- Phase 1 – external brute‑force -----
	if not self._pivot_done:
	if self._external_failed < self.max_failures:
	logs.append(
	LogEntry(
	timestamp=env._rand_timestamp(),
	source="ssh",
	severity="WARN",
	message=f"Failed password for invalid user admin from {self.attacker_ip}",
	ip=self.attacker_ip,
	user="admin",
	hostname="gateway.corp",
	event_type="login",
	)
	)
	self._external_failed += 1
	return logs
	# success – pivot host compromised
	logs.append(
	LogEntry(
	timestamp=env._rand_timestamp(),
	source="ssh",
	severity="INFO",
	message=f"Accepted password for svc_user from {self.attacker_ip} on {self.pivot_host}",
	ip=self.attacker_ip,
	user="svc_user",
	hostname=self.pivot_host,
	event_type="login",
	)
	)
	env._state.compromised_hosts.append(self.pivot_host)
	self._pivot_done = True
	return logs

	# ----- Phase 2 – internal SSH hop -----
	if self._pivot_done and not self._internal_done:
	logs.append(
	LogEntry(
	timestamp=env._rand_timestamp(),
	source="ssh",
	severity="INFO",
	message=f"{self.pivot_host} initiated SSH connection to {self.internal_host}",
	ip=self.pivot_host,
	user="svc_user",
	hostname=self.internal_host,
	event_type="network_conn",
	)
	)
	# after a couple of steps the internal host gets admin rights
	if self._internal_failed >= 2:
	logs.append(
	LogEntry(
	timestamp=env._rand_timestamp(),
	source="ssh",
	severity="INFO",
	message=f"New admin account created on {self.internal_host} by svc_user",
	ip=self.internal_host,
	user="admin",
	hostname=self.internal_host,
	event_type="login",
	)
	)
	env._state.compromised_hosts.append(self.internal_host)
	self._internal_done = True
	env._state.data_exfiltrated = True
	else:
	self._internal_failed += 1
	return logs

	# Phase 3 – attack already succeeded – nothing extra to emit
	return logs

	def is_mitigated(self, env: AIGymEnv) -> bool:
	compromised = set(env._state.compromised_hosts)
	isolated = set(env._isolated_hosts)
	return len(compromised & isolated) > 0


	# ----------------------------------------------------------------------
	# Task 3 – APT Multi‑Stage (hard difficulty)
	# ----------------------------------------------------------------------
	@dataclass
	class APTMultiStageTask(BaseTask):
	"""Hard scenario — a full Advanced Persistent Threat kill-chain.

	Phases:
	1️⃣ Phishing — Spear‑phishing email delivers a malicious document
	to a user workstation. The logs show an inbound email with a
	suspicious attachment, followed by an outbound C2 beacon.
	2️⃣ Credential Theft — The compromised workstation runs a
	credential‑dumping tool (Mimikatz‑style). Logs show LSASS
	access and new logon events with harvested service accounts.
	3️⃣ Lateral Movement — Stolen creds are used to hop through
	two internal hosts. Logs show WinRM / SSH connections between
	internal machines.
	4️⃣ Exfiltration — Attacker stages data on a staging host and
	tunnels it out via DNS over HTTPS. Logs show large DNS‑TXT
	responses and unusual outbound traffic volume.

	Mitigation requires the analyst to:
	• Isolate the phished workstation AND
	• Block the C2 IP
	Partial credit is given for blocking C2 alone or isolating any
	compromised host.
	"""

	# ---------- configurable knobs ----------
	c2_ip: str = "198.51.100.23"
	phish_target_user: str = "jdoe"
	phished_host: str = "ws-PC042.corp"
	staging_host: str = "file-srv-01.corp"
	dc_host: str = "dc-01.corp"
	max_phish_steps: int = 2 # steps before the user clicks the link

	# ---------- internal counters ----------
	_phase: int = field(default=0, init=False) # 0–4
	_phase_counter: int = field(default=0, init=False) # steps within phase
	_user_clicked: bool = field(default=False, init=False)
	_creds_dumped: bool = field(default=False, init=False)
	_lateral_hops: int = field(default=0, init=False)
	_exfil_started: bool = field(default=False, init=False)
	hosts_investigated: Set[str] = field(default_factory=set, init=False)

	def initialize(self, env: AIGymEnv) -> None:
	env._state.stage = "initial_access"
	env._state.attacker_ips = [self.c2_ip]
	env._state.pivot_host = self.phished_host
	env._state.internal_host = self.staging_host
	self._phase = 0
	self._phase_counter = 0
	self._user_clicked = False
	self._creds_dumped = False
	self._lateral_hops = 0
	self._exfil_started = False
	self.hosts_investigated = set()

	# ------------------------------------------------------------------
	def advance(self, env: AIGymEnv) -> List[LogEntry]:
	logs: List[LogEntry] = []
	self._phase_counter += 1

	# ---- Phase 0 : Phishing delivery ----
	if self._phase == 0:
	logs.append(
	LogEntry(
	timestamp=env._rand_timestamp(),
	source="email",
	severity="INFO",
	message=(
	f"Inbound email to {self.phish_target_user}@corp.local "
	f"from hr-benefits@secure‑docs.com "
	f"subject='Q1 Compensation Update' attachment=Q1_Update.xlsm"
	),
	ip=self.c2_ip,
	user=self.phish_target_user,
	hostname=self.phished_host,
	event_type="email",
	)
	)
	if self._phase_counter >= self.max_phish_steps:
	# User clicks the doc — macro fires
	logs.append(
	LogEntry(
	timestamp=env._rand_timestamp(),
	source="endpoint",
	severity="WARN",
	message=(
	f"EXCEL.EXE spawned powershell.exe on {self.phished_host} "
	f"(user={self.phish_target_user})"
	),
	ip=None,
	user=self.phish_target_user,
	hostname=self.phished_host,
	event_type="process",
	)
	)
	logs.append(
	LogEntry(
	timestamp=env._rand_timestamp(),
	source="firewall",
	severity="WARN",
	message=(
	f"Outbound HTTPS connection from {self.phished_host} "
	f"to {self.c2_ip}:443 (TLS SNI: cdn‑assets.com)"
	),
	ip=self.c2_ip,
	user=self.phish_target_user,
	hostname=self.phished_host,
	event_type="network_conn",
	)
	)
	self._user_clicked = True
	env._state.compromised_hosts.append(self.phished_host)
	self._phase = 1
	self._phase_counter = 0
	return logs

	# ---- Phase 1 : Credential theft ----
	if self._phase == 1:
	# First step in this phase — LSASS access
	if not self._creds_dumped:
	logs.append(
	LogEntry(
	timestamp=env._rand_timestamp(),
	source="endpoint",
	severity="CRITICAL",
	message=(
	f"Suspicious LSASS memory access on {self.phished_host} "
	f"by process rundll32.exe (PID 4728)"
	),
	ip=None,
	user="SYSTEM",
	hostname=self.phished_host,
	event_type="process",
	)
	)
	logs.append(
	LogEntry(
	timestamp=env._rand_timestamp(),
	source="auth",
	severity="WARN",
	message=(
	f"New logon: svc_backup@corp.local on {self.phished_host} "
	f"(logon type 9 — NewCredentials)"
	),
	ip=None,
	user="svc_backup",
	hostname=self.phished_host,
	event_type="login",
	)
	)
	self._creds_dumped = True
	env._state.stage = "lateral_movement"
	self._phase = 2
	self._phase_counter = 0
	return logs

	# ---- Phase 2 : Lateral movement (2 hops) ----
	if self._phase == 2:
	if self._lateral_hops == 0:
	# Hop 1: phished_host → dc_host
	logs.append(
	LogEntry(
	timestamp=env._rand_timestamp(),
	source="auth",
	severity="WARN",
	message=(
	f"Kerberos TGS request for host/{self.dc_host} "
	f"from {self.phished_host} (user=svc_backup)"
	),
	ip=self.phished_host,
	user="svc_backup",
	hostname=self.dc_host,
	event_type="login",
	)
	)
	logs.append(
	LogEntry(
	timestamp=env._rand_timestamp(),
	source="ssh",
	severity="INFO",
	message=(
	f"WinRM session opened on {self.dc_host} "
	f"by svc_backup from {self.phished_host}"
	),
	ip=self.phished_host,
	user="svc_backup",
	hostname=self.dc_host,
	event_type="network_conn",
	)
	)
	env._state.compromised_hosts.append(self.dc_host)
	self._lateral_hops = 1
	return logs

	if self._lateral_hops == 1:
	# Hop 2: dc_host → staging_host
	logs.append(
	LogEntry(
	timestamp=env._rand_timestamp(),
	source="ssh",
	severity="INFO",
	message=(
	f"{self.dc_host} initiated SMB connection to "
	f"{self.staging_host} (share: \\\\{self.staging_host}\\data$)"
	),
	ip=self.dc_host,
	user="svc_backup",
	hostname=self.staging_host,
	event_type="network_conn",
	)
	)
	logs.append(
	LogEntry(
	timestamp=env._rand_timestamp(),
	source="endpoint",
	severity="WARN",
	message=(
	f"Batch file copy: 847 files (2.3 GB) staged in "
	f"C:\\Windows\\Temp\\~cache on {self.staging_host}"
	),
	ip=None,
	user="svc_backup",
	hostname=self.staging_host,
	event_type="file_access",
	)
	)
	env._state.compromised_hosts.append(self.staging_host)
	self._lateral_hops = 2
	env._state.stage = "exfiltration"
	self._phase = 3
	self._phase_counter = 0
	return logs

	return logs

	# ---- Phase 3 : Exfiltration via DNS tunneling ----
	if self._phase == 3:
	if not self._exfil_started:
	logs.append(
	LogEntry(
	timestamp=env._rand_timestamp(),
	source="dns",
	severity="WARN",
	message=(
	f"Anomalous DNS-over-HTTPS traffic from {self.staging_host} "
	f"to {self.c2_ip} — 340 TXT queries in 60 s "
	f"(avg payload 220 bytes)"
	),
	ip=self.c2_ip,
	user=None,
	hostname=self.staging_host,
	event_type="network_conn",
	)
	)
	logs.append(
	LogEntry(
	timestamp=env._rand_timestamp(),
	source="firewall",
	severity="CRITICAL",
	message=(
	f"High-volume outbound: {self.staging_host} → "
	f"{self.c2_ip}:443 \| 2.1 GB transferred in 180 s"
	),
	ip=self.c2_ip,
	user=None,
	hostname=self.staging_host,
	event_type="network_conn",
	)
	)
	self._exfil_started = True

	if self._phase_counter >= 2:
	env._state.data_exfiltrated = True

	return logs

	return logs

	# ------------------------------------------------------------------
	def is_mitigated(self, env: AIGymEnv) -> bool:
	"""Full mitigation requires BOTH:
	1. The C2 IP is blocked OR the phished host is isolated.
	2. At least one compromised host is isolated.
	"""
	c2_blocked = self.c2_ip in env._blocked_ips
	phished_isolated = self.phished_host in env._isolated_hosts
	any_compromised_isolated = bool(
	set(env._state.compromised_hosts) & set(env._isolated_hosts)
	)
	return (c2_blocked or phished_isolated) and any_compromised_isolated


	# ----------------------------------------------------------------------
	# Exported symbols for `import *`
	# ----------------------------------------------------------------------
	__all__ = [
	"BaseTask",
	"BruteForceSSHTask",
	"LateralMovementTask",
	"APTMultiStageTask",
	]