bolyai/CIS490
3
0
Fork 0
Code Issues 5 Pull requests Projects 1 Releases Packages Wiki Activity Actions

Driver v2: sample-profile-driven workloads (Tier-2 + Tier-3)

Browse source 
The v1 driver ran ``yes > /dev/null`` for every sample, which
produced the same envelope shape regardless of which malware family
the orchestrator claimed to be running. That's a poor training
signal: the model sees identical /proc + QMP traces tagged
"cryptominer" / "ransomware" / "RAT" with no distinguishing
features. v2 fixes this.

What landed:

  exploits/workloads.py — six ``Workload`` profiles, each producing
    a distinct in-session shell command pair (start_cmd / stop_cmd)
    that backgrounds a profile-shaped loop:

      cpu-saturate    — sustained 1-vCPU saturation (XMRig shape)
      scan-and-dial   — periodic SYN-style probes across 10.200.0.0/24
                        + dial-home to gateway (Mirai shape)
      io-walk         — fs traversal + 4 KiB urandom writes, periodic
                        re-read (ransomware shape)
      bursty-c2       — long idle, periodic 3-packet TCP egress burst
                        (Dridex C2 beacon shape)
      low-and-slow    — minimal CPU + periodic awk-driven memory churn
                        (Kovter / fileless shape)
      shell-resident  — single long-lived TCP socket pinned to gateway
                        with periodic 6-byte command ticks (RAT shape)

  Each profile uses a /tmp/.cis490-workload-<profile>.{pid,sh} pair so
  the stop_cmd can cleanly kill the loop and its descendants.

  exploits/driver.py — MSFExploitDriver now accepts an optional
    ``Sample``. With one supplied, ``infected_running`` dispatches to
    the matching workload via exploits.workloads.workload_for(); the
    ``sample_executed`` event records profile + sample name + sample
    kind so the trainer can join cleanly. Without a sample, the v1
    yes-loop path remains unchanged (backwards compat).

  tools/vm_load_controller.py — the same dispatch on the Tier-2 path
    (no exploit, real Alpine guest driven over the serial console).
    A fleet wave now produces six visually distinct envelopes per
    wave whether the underlying mode is Tier 2 or Tier 3.

  tools/run_real_vm_demo.py — accepts ``--sample <name>`` (or
    SAMPLE_NAME env from the fleet runner) + auto-wires QMP + agent
    sockets into the EpisodeConfig so all three new collectors
    (sources 2, 4, 5) run alongside source 1 by default.

  tools/run_tier3_demo.py — same ``--sample`` plumbing for the
    exploit-driven path.

Tests: 86 pass (was 82). New v2 cases:
  - profile dispatch routes infected_running to the workload's
    start_cmd (NOT the v1 yes-loop) when a Sample is set
  - all six profiles produce distinct start_cmds (the property the
    ML model needs)
  - unknown profile string falls back to cpu-saturate with a warning
  - v1 path (no Sample) still uses yes-loop (backwards compat)

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
max 2026-04-30 00:06:15 -05:00
parent 1b6c7b2f4a
commit b80986d99c
6 changed files with 474 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

73
exploits/driver.py
View file

@ -31,8 +31,11 @@ import time
from dataclasses import dataclass
from typing import Callable
from samples.manifest import Sample
from .modules import ModuleConfig
from .msfrpc import MSFRpcClient, wait_for_new_session
from .workloads import Workload, workload_for
log = logging.getLogger("cis490.exploits.driver")
@ -44,17 +47,23 @@ EmitEvent = Callable[..., None]
class DriverConfig:
target_ip: str
session_open_timeout_s: float = 30.0
# Workload command used to mimic XMRig-class infected_running shape
# in a real session. Kept simple on purpose — anything observable
# from outside the guest works for the dataset; we'll drop in a
# real sample at Tier 4.
# Driver v1 fallback workload — used only when no Sample is passed
# in (Sample-driven runs override these via exploits.workloads).
# We keep the v1 path so existing callers keep working unchanged.
workload_cmd: str = "yes > /dev/null"
# How we kill the workload at dormant time.
workload_kill_cmd: str = "pkill yes; true"
class MSFExploitDriver:
"""Phase-to-msfrpc adapter. One instance per episode."""
"""Phase-to-msfrpc adapter. One instance per episode.
When constructed with a ``Sample``, the driver dispatches the
``infected_running`` / ``dormant`` workload through
``exploits.workloads`` so the in-session behaviour matches the
sample's profile (cpu-saturate, scan-and-dial, io-walk, bursty-c2,
low-and-slow, shell-resident). Without a sample, falls back to
the v1 single-command workload — useful for the very first
Tier-3 smoke runs."""
def __init__(
self,
@ -62,11 +71,15 @@ class MSFExploitDriver:
module: ModuleConfig,
cfg: DriverConfig,
emit_event: EmitEvent,
*,
sample: Sample | None = None,
) -> None:
self.client = client
self.module = module
self.cfg = cfg
self.emit = emit_event
self.sample = sample
self.workload: Workload | None = workload_for(sample)
self._sessions_seen_at_arm: set[int] = set()
self._session_id: int | None = None
@ -86,6 +99,9 @@ class MSFExploitDriver:
payload=self.module.payload_path,
target_ip=self.cfg.target_ip,
preexisting_sessions=sorted(self._sessions_seen_at_arm),
sample=self.sample.name if self.sample else None,
sample_kind=self.sample.kind if self.sample else None,
workload_profile=self.workload.profile if self.workload else None,
)
def teardown(self) -> None:
@ -178,24 +194,43 @@ class MSFExploitDriver:
if self._session_id is None:
log.warning("infected_running with no session — skipping workload")
return
self.client.session_shell_write(
self._session_id,
f"nohup sh -c {_shquote(self.cfg.workload_cmd)} </dev/null "
f">/dev/null 2>&1 & disown",
)
self.emit(
"sample_executed",
session_id=self._session_id,
command=self.cfg.workload_cmd,
)
if self.workload is not None:
# Driver v2 — profile-matched workload.
self.client.session_shell_write(self._session_id, self.workload.start_cmd)
self.emit(
"sample_executed",
session_id=self._session_id,
profile=self.workload.profile,
description=self.workload.description,
sample=self.sample.name if self.sample else None,
)
else:
# Driver v1 fallback.
self.client.session_shell_write(
self._session_id,
f"nohup sh -c {_shquote(self.cfg.workload_cmd)} </dev/null "
f">/dev/null 2>&1 & disown",
)
self.emit(
"sample_executed",
session_id=self._session_id,
command=self.cfg.workload_cmd,
)
def _stop_workload(self) -> None:
if self._session_id is None:
return
self.client.session_shell_write(
self._session_id, self.cfg.workload_kill_cmd,
if self.workload is not None:
self.client.session_shell_write(self._session_id, self.workload.stop_cmd)
else:
self.client.session_shell_write(
self._session_id, self.cfg.workload_kill_cmd,
)
self.emit(
"session_dormant",
session_id=self._session_id,
profile=self.workload.profile if self.workload else None,
)
self.emit("session_dormant", session_id=self._session_id)
def _shquote(s: str) -> str:

235
exploits/workloads.py Normal file
View file

@ -0,0 +1,235 @@
"""Per-sample-profile post-exploit workloads (driver v2).
The Tier-3 driver lands a session and then needs to drive *something*
in that session for the ``infected_running`` phase. Driver v1 ran
``yes > /dev/null`` for every sample, which is fine for proving the
pipe but is the wrong shape for ML — every Tier-3 episode produces
the same envelope regardless of which malware family we said it was.
Driver v2 maps ``sample.profile`` from the manifest to a distinct
in-session workload so each profile's envelope is observably
different on every collector:
cpu-saturate → 1-vCPU saturation, very low IO/net (XMRig shape)
scan-and-dial → SYN scans across the bridge IP space + periodic
dial-home (Mirai shape)
io-walk → fs traversal + random write spikes (ransomware shape)
bursty-c2 → long idle, periodic short TCP egress bursts (Dridex)
low-and-slow → minimal CPU, periodic memory churn (Kovter)
shell-resident → one long-lived TCP socket pinned to a bridge IP,
occasional small command bursts (RAT)
Each profile returns a small shell command that backgrounds a loop
inside the session. The driver can stop them by killing the loop's
PID file or via a profile-specific kill command.
This module is intentionally *behaviorally diverse but harmless* —
it does NOT execute real malware. Real binaries land via the Tier-4
fetch+run path (separate work). What this gives us today is six
distinguishable in-guest envelopes the ML model can learn to
discriminate between *and* fall back to when a real sample isn't yet
staged.
"""
from __future__ import annotations
import logging
from dataclasses import dataclass
from samples.manifest import Sample
log = logging.getLogger("cis490.exploits.workloads")
@dataclass(frozen=True)
class Workload:
"""A pair of shell commands executable in a Metasploit shell session.
``start_cmd`` backgrounds a loop and writes its PID to ``pid_path``.
``stop_cmd`` kills the loop using that PID file. Both commands are
expected to be POSIX-shell compatible and to leave the session in
a usable state on completion (return code 0 on the prompt)."""
profile: str
start_cmd: str
stop_cmd: str
description: str
@property
def pid_path(self) -> str:
return f"/tmp/.cis490-workload-{self.profile}.pid"
def _wrap_loop(name: str, body: str) -> Workload:
"""Common pattern: write a small wrapper script that loops ``body``,
background it, and stash the wrapper's PID. Stop kills that PID +
its child group."""
pid_path = f"/tmp/.cis490-workload-{name}.pid"
script_path = f"/tmp/.cis490-workload-{name}.sh"
# Triple-quote the body into a heredoc so single-quotes inside the
# body don't conflict with our outer single-quoting.
start = (
f"cat > {script_path} <<'CIS490_EOF'\n"
f"#!/bin/sh\n"
f"trap 'exit 0' TERM INT\n"
f"while :; do\n"
f"{body}\n"
f"done\n"
f"CIS490_EOF\n"
f"chmod +x {script_path}; "
f"nohup sh {script_path} </dev/null >/dev/null 2>&1 &\n"
f"echo $! > {pid_path}\n"
f"disown\n"
)
stop = (
f"if [ -f {pid_path} ]; then "
f" kill -- -$(cat {pid_path}) 2>/dev/null; "
f" kill $(cat {pid_path}) 2>/dev/null; "
f" rm -f {pid_path} {script_path}; "
f"fi; true\n"
)
return Workload(profile=name, start_cmd=start, stop_cmd=stop,
description="(generated)")
# ---------------------------------------------------------------------------
# Profile factories — each returns a Workload tuned to that family
# ---------------------------------------------------------------------------
def _cpu_saturate() -> Workload:
"""XMRig-class — sustained single-vCPU saturation, no IO, no net."""
body = " yes > /dev/null 2>&1 &\n wait $!\n"
w = _wrap_loop("cpu-saturate", body)
return Workload(
profile="cpu-saturate",
start_cmd=w.start_cmd,
stop_cmd=w.stop_cmd,
description="100% CPU on 1 vCPU; no IO, no net",
)
def _scan_and_dial() -> Workload:
"""Mirai-class — TCP SYN-style probe of bridge subnet + occasional
"dial home" to the gateway. Heavy net, moderate CPU."""
body = (
" for i in 1 2 3 4 5 6 7 8 9 10; do\n"
" (echo > /dev/tcp/10.200.0.$((i+1))/23) 2>/dev/null &\n"
" (echo > /dev/tcp/10.200.0.$((i+1))/2323) 2>/dev/null &\n"
" done\n"
" wait\n"
" (echo dial-home > /dev/tcp/10.200.0.1/4444) 2>/dev/null\n"
" sleep 2\n"
)
w = _wrap_loop("scan-and-dial", body)
return Workload(
profile="scan-and-dial",
start_cmd=w.start_cmd,
stop_cmd=w.stop_cmd,
description="Periodic SYN-style scan across bridge IPs + dial-home",
)
def _io_walk() -> Workload:
"""Cryptolocker-class — fs traversal + write spikes. Heavy disk."""
body = (
" mkdir -p /tmp/.cis490-victim\n"
" for n in 1 2 3 4 5 6 7 8; do\n"
" dd if=/dev/urandom of=/tmp/.cis490-victim/f$n bs=4k count=64 2>/dev/null\n"
" done\n"
" for f in /tmp/.cis490-victim/*; do cat $f > /dev/null; done\n"
" sleep 1\n"
)
w = _wrap_loop("io-walk", body)
return Workload(
profile="io-walk",
start_cmd=w.start_cmd,
stop_cmd=w.stop_cmd,
description="FS traversal + random-data writes, periodic re-read",
)
def _bursty_c2() -> Workload:
"""Dridex-class — long idle, periodic small TCP burst to a fixed
peer (the bridge gateway)."""
body = (
" sleep 25\n"
" for i in 1 2 3; do\n"
" (echo c2-beacon-$$-$i > /dev/tcp/10.200.0.1/4445) 2>/dev/null\n"
" sleep 1\n"
" done\n"
)
w = _wrap_loop("bursty-c2", body)
return Workload(
profile="bursty-c2",
start_cmd=w.start_cmd,
stop_cmd=w.stop_cmd,
description="Long idle + periodic 3-packet egress burst to gateway",
)
def _low_and_slow() -> Workload:
"""Kovter-class — low CPU, periodic memory churn, no on-disk
artifact. The hardest envelope to label from /proc alone."""
body = (
" sleep 8\n"
" awk 'BEGIN { for(i=0;i<200000;i++) a[i]=i*i; }' >/dev/null 2>&1\n"
" sleep 4\n"
)
w = _wrap_loop("low-and-slow", body)
return Workload(
profile="low-and-slow",
start_cmd=w.start_cmd,
stop_cmd=w.stop_cmd,
description="Periodic memory churn (~200k array allocs) on a slow cycle",
)
def _shell_resident() -> Workload:
"""RAT-style — keep a single TCP socket open to the gateway with
occasional command bursts. Long-lived flow, small bytes."""
# nc on Metasploitable2 is GNU netcat; on busybox it's also there.
# We use plain bash /dev/tcp redirects to avoid depending on nc.
body = (
" exec 3<>/dev/tcp/10.200.0.1/4446 2>/dev/null && {\n"
" for i in 1 2 3 4 5 6; do\n"
" echo cmd-tick-$i >&3\n"
" sleep 5\n"
" done\n"
" exec 3<&-; exec 3>&-\n"
" }\n"
" sleep 5\n"
)
w = _wrap_loop("shell-resident", body)
return Workload(
profile="shell-resident",
start_cmd=w.start_cmd,
stop_cmd=w.stop_cmd,
description="Resident TCP connection to gateway with periodic ticks",
)
_FACTORIES = {
"cpu-saturate": _cpu_saturate,
"scan-and-dial": _scan_and_dial,
"io-walk": _io_walk,
"bursty-c2": _bursty_c2,
"low-and-slow": _low_and_slow,
"shell-resident": _shell_resident,
}
def workload_for(sample: Sample | None) -> Workload | None:
"""Return the Workload matching ``sample.profile``, or None when
no sample is supplied (driver v1 fallback path)."""
if sample is None:
return None
factory = _FACTORIES.get(sample.profile)
if factory is None:
log.warning("no workload profile for %r; falling back to cpu-saturate", sample.profile)
return _cpu_saturate()
return factory()
def all_profiles() -> list[str]:
return sorted(_FACTORIES.keys())

105
tests/test_exploits.py
View file

@ -252,6 +252,111 @@ def test_teardown_kills_session_and_logs_out() -> None:
# Driver wired into a real EpisodeRunner — events land in events.jsonl
# -----------------------------------------------------------------------
# -----------------------------------------------------------------------
# Driver v2 — sample-profile-driven workloads
# -----------------------------------------------------------------------
def test_v2_uses_profile_workload_for_cpu_saturate() -> None:
"""When constructed with a Sample, the driver should send the
profile's start_cmd at infected_running rather than the v1
yes-loop. The actual command body is owned by exploits.workloads
and tested there; here we just confirm dispatch."""
from samples.manifest import Sample as _Sample
cfg = load_module_config(MODULES_DIR / "vsftpd_234_backdoor.toml")
client = FakeMSFRpcClient(
sessions_after_fire={1: {"type": "shell", "tunnel_peer": "x:21"}},
)
events: list[tuple[str, dict]] = []
sample = _Sample(
name="xmrig-cryptominer",
family="XMRig",
category="cryptominer",
profile="cpu-saturate",
)
driver = MSFExploitDriver(
client=client, # type: ignore[arg-type]
module=cfg,
cfg=DriverConfig(target_ip="10.200.0.10", session_open_timeout_s=0.5),
emit_event=lambda ev, **kw: events.append((ev, kw)),
sample=sample,
)
driver.setup()
driver.set_phase("armed")
driver.set_phase("infecting")
driver.set_phase("infected_running")
driver.set_phase("dormant")
driver.teardown()
# The shell command sent at infected_running should be the
# profile's multi-line wrapper — NOT the v1 single-yes line.
starts = [w for (_, w) in client.shell_writes if "yes > /dev/null" in w and "cis490-workload" not in w]
assert starts == [], "v2 driver must not send the v1 yes-loop when a Sample is supplied"
# The driver_setup event records sample + workload metadata.
setup_events = [kw for (e, kw) in events if e == "driver_setup"]
assert setup_events
assert setup_events[0]["sample"] == "xmrig-cryptominer"
assert setup_events[0]["sample_kind"] == "mimic"
assert setup_events[0]["workload_profile"] == "cpu-saturate"
# sample_executed carries the profile name + description.
se = [kw for (e, kw) in events if e == "sample_executed"]
assert se
assert se[0]["profile"] == "cpu-saturate"
assert se[0]["sample"] == "xmrig-cryptominer"
def test_v2_distinct_workloads_per_profile() -> None:
"""Two different profiles must produce *different* shell commands.
This is the property that gives the ML model varied envelopes to
learn from."""
from exploits.workloads import all_profiles, workload_for
from samples.manifest import Sample as _Sample
profiles = all_profiles()
assert len(profiles) >= 4
seen_starts: set[str] = set()
for p in profiles:
s = _Sample(name=f"x-{p}", family="X", category="rat", profile=p)
w = workload_for(s)
assert w is not None
seen_starts.add(w.start_cmd)
# Every profile must have a distinct start_cmd.
assert len(seen_starts) == len(profiles), \
"two profiles produced the same workload — ML diversity is at risk"
def test_v2_unknown_profile_falls_back_to_cpu_saturate() -> None:
from exploits.workloads import workload_for
from samples.manifest import Sample as _Sample
s = _Sample(name="weird", family="X", category="rat", profile="not-a-real-profile")
w = workload_for(s)
assert w is not None
assert w.profile == "cpu-saturate"
def test_v1_path_still_works_when_no_sample() -> None:
"""Ensure backwards compat: a driver constructed without a sample
uses the original yes-loop workload."""
cfg = load_module_config(MODULES_DIR / "vsftpd_234_backdoor.toml")
client = FakeMSFRpcClient(sessions_after_fire={1: {"type": "shell"}})
driver = MSFExploitDriver(
client=client, # type: ignore[arg-type]
module=cfg,
cfg=DriverConfig(target_ip="10.200.0.10", session_open_timeout_s=0.5),
emit_event=lambda *a, **kw: None,
)
driver.setup()
driver.set_phase("armed")
driver.set_phase("infecting")
driver.set_phase("infected_running")
driver.teardown()
assert any("yes > /dev/null" in w for (_, w) in client.shell_writes)
def test_driver_events_persist_to_events_jsonl(tmp_path: Path) -> None:
"""When the driver is connected to a real EpisodeRunner, the
events it emits must show up in the episode's events.jsonl with

28
tools/run_real_vm_demo.py
View file

@ -28,6 +28,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
sys.path.insert(0, str(Path(__file__).resolve().parent))
from orchestrator.episode import EpisodeConfig, EpisodeRunner # noqa: E402
from samples.manifest import SampleManifest # noqa: E402
from vm_load_controller import VMLoadController # noqa: E402
from vm_serial import SerialClient # noqa: E402
@ -83,6 +84,16 @@ def main() -> int:
default=120.0,
help="how long to wait for serial login prompt",
)
parser.add_argument(
"--sample",
default=os.environ.get("SAMPLE_NAME"),
help="Pick a workload profile from the manifest by name. Fleet runner "
"passes this via SAMPLE_NAME env. If unset, runs the v1 yes-loop.",
)
parser.add_argument(
"--manifest",
default=str(Path(__file__).resolve().parent.parent / "samples" / "manifest.toml"),
)
args = parser.parse_args()
logging.basicConfig(
@ -93,6 +104,17 @@ def main() -> int:
repo_root = Path(__file__).resolve().parent.parent
launcher = repo_root / "vm" / "launch_demo.sh"
# Resolve sample if requested.
sample = None
if args.sample:
manifest = SampleManifest.load(args.manifest)
sample = next((s for s in manifest.samples if s.name == args.sample), None)
if sample is None:
log.error("sample %r not in manifest %s", args.sample, args.manifest)
return 2
log.info("using sample=%s profile=%s kind=%s",
sample.name, sample.profile, sample.kind)
run_dir = Path(args.run_dir)
# Wipe any stale sockets/pidfile from a previous run.
if run_dir.exists():
@ -137,9 +159,11 @@ def main() -> int:
serial.connect()
serial.login(boot_timeout_s=args.boot_timeout)
controller = VMLoadController(serial)
controller = VMLoadController(serial, sample=sample)
controller.setup()
qmp_sock = run_dir / "qmp.sock"
agent_sock = run_dir / "agent.sock"
cfg = EpisodeConfig(
target_pid=qemu_pid,
duration_s=sum(d for _, d in DEFAULT_SCHEDULE),
@ -148,6 +172,8 @@ def main() -> int:
phase_schedule=DEFAULT_SCHEDULE,
image_name="alpine-3.21-cloudinit",
snapshot_name="baseline-v1",
qmp_socket=qmp_sock if qmp_sock.exists() else None,
guest_agent_socket=agent_sock if agent_sock.exists() else None,
)
result = EpisodeRunner(cfg, on_phase=controller.set_phase).run()

22
tools/run_tier3_demo.py
View file

@ -37,6 +37,7 @@ from exploits.driver import DriverConfig, MSFExploitDriver # noqa: E402
from exploits.modules import load_module_config # noqa: E402
from exploits.msfrpc import MSFRpcClient, MSFRpcConfig # noqa: E402
from orchestrator.episode import EpisodeConfig, EpisodeRunner # noqa: E402
from samples.manifest import SampleManifest # noqa: E402
# Same envelope shape as Tier 2 so plots are comparable. Slightly more
@ -128,6 +129,16 @@ def main() -> int:
default=180.0,
help="how long to wait for the guest's vulnerable service to listen",
)
parser.add_argument(
"--sample",
default=os.environ.get("SAMPLE_NAME"),
help="Pick a workload profile from the manifest by name. Fleet runner "
"passes this via SAMPLE_NAME env. Without it, falls back to the v1 yes-loop.",
)
parser.add_argument(
"--manifest",
default=str(Path(__file__).resolve().parent.parent / "samples" / "manifest.toml"),
)
args = parser.parse_args()
logging.basicConfig(
@ -152,6 +163,16 @@ def main() -> int:
module = load_module_config(module_path)
log.info("module loaded: %s (%s)", module.name, module.module_path)
sample = None
if args.sample:
manifest = SampleManifest.load(args.manifest)
sample = next((s for s in manifest.samples if s.name == args.sample), None)
if sample is None:
log.error("sample %r not in manifest %s", args.sample, args.manifest)
return 2
log.info("sample=%s profile=%s kind=%s",
sample.name, sample.profile, sample.kind)
run_dir = Path(args.run_dir)
if run_dir.exists():
import shutil
@ -205,6 +226,7 @@ def main() -> int:
module=module,
cfg=DriverConfig(target_ip=args.target_ip),
emit_event=runner.emit_event,
sample=sample,
)
runner.on_phase = driver.set_phase

38
tools/vm_load_controller.py
View file

@ -22,16 +22,34 @@ fire and a real sample.
from __future__ import annotations
import logging
import sys
from pathlib import Path
from vm_serial import SerialClient
# Allow running as a script (sibling of tools/).
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from exploits.workloads import Workload, workload_for # noqa: E402
from samples.manifest import Sample # noqa: E402
log = logging.getLogger("cis490.vm_load_controller")
class VMLoadController:
def __init__(self, serial: SerialClient) -> None:
"""Drives a real Alpine guest through the phase schedule for
Tier 2 (no exploit). Workload is chosen by ``sample.profile`` —
same profile catalog as the Tier-3 driver so a fleet wave
produces matched envelopes whether or not an exploit fires.
Without a sample, falls back to the original cpu-saturate yes-loop
(the original Tier-2 demo behaviour)."""
def __init__(self, serial: SerialClient, sample: Sample | None = None) -> None:
self.s = serial
self.sample = sample
self.workload: Workload | None = workload_for(sample)
def setup(self) -> None:
# Kill any pre-existing load and clear scratch space.
@ -44,7 +62,8 @@ class VMLoadController:
# ---- phases ---------------------------------------------------------
def set_phase(self, phase: str) -> None:
log.info("vm phase -> %s", phase)
log.info("vm phase -> %s (profile=%s)",
phase, self.workload.profile if self.workload else "v1")
if phase == "clean":
self._kill_load()
elif phase == "armed":
@ -56,10 +75,12 @@ class VMLoadController:
)
elif phase == "infected_running":
self._kill_load()
# Background CPU burner. `nohup` + `&` + redirects to detach.
self.s.run(
"nohup sh -c 'yes > /dev/null' </dev/null >/dev/null 2>&1 & disown"
)
if self.workload is not None:
self.s.run(self.workload.start_cmd)
else:
self.s.run(
"nohup sh -c 'yes > /dev/null' </dev/null >/dev/null 2>&1 & disown"
)
elif phase == "dormant":
self._kill_load()
else:
@ -68,5 +89,8 @@ class VMLoadController:
# ---- internals ------------------------------------------------------
def _kill_load(self) -> None:
# `true` at the end so the run() exit status is always 0.
if self.workload is not None:
self.s.run(self.workload.stop_cmd)
# Always sweep the v1 leftover commands too, in case we just
# switched profiles mid-fleet-run.
self.s.run("pkill yes 2>/dev/null; pkill stress-ng 2>/dev/null; true")