disk_panic/server/scenarios.py

"""The three task scenarios. Each builds a fresh VFS in a specific broken state."""
from __future__ import annotations

import hashlib
from typing import Dict, Tuple

from disk_panic.server.vfs import VFS, VFSFile

# Sentinel content for the protected audit directory. Hash is captured at reset.
_AUDIT_FILES = {
    "/var/log/audit/audit.log": b"type=LOGIN audit_id=001 user=root\n" * 64,
    "/var/log/audit/audit.log.1": b"type=LOGIN audit_id=000 user=root\n" * 128,
}


def _base_fs() -> VFS:
    vfs = VFS()
    # /etc skeleton
    vfs.files["/etc/hostname"] = VFSFile(b"prod-app-01\n", 12)
    vfs.files["/etc/os-release"] = VFSFile(b"NAME=Linux\n", 11)
    # /var/log with audit (protected) and some normal logs
    for path, content in _AUDIT_FILES.items():
        vfs.files[path] = VFSFile(content=content, size=len(content), protected=True)
    vfs.files["/var/log/syslog"] = VFSFile(b"boot ok\n", 256 * 1024)
    vfs.files["/var/log/nginx/access.log"] = VFSFile(b"", 128 * 1024 * 1024)  # 128M
    # /app binary
    vfs.files["/opt/app/app.bin"] = VFSFile(b"#!/bin/app\n", 40 * 1024 * 1024)
    return vfs


def build_easy() -> Tuple[VFS, Dict[str, object]]:
    vfs = _base_fs()
    # One huge rotated log is nearly filling the disk.
    vfs.files["/var/log/nginx/access.log.1.gz"] = VFSFile(
        content=b"\x1f\x8b" + b"X" * 1024,
        size=int(8.7 * 1024 ** 3),  # 8.7G
    )
    vfs.services["app"] = "active"
    targets = {
        "audit_sha": _audit_sha(vfs),
        "requires_service_active": False,
        "requires_logrotate": False,
    }
    return vfs, targets


def build_medium() -> Tuple[VFS, Dict[str, object]]:
    vfs = _base_fs()
    vfs.files["/var/log/journal/system.journal"] = VFSFile(
        content=b"", size=int(8.5 * 1024 ** 3)
    )
    # Service failed because disk was full
    vfs.services["app"] = "failed"
    targets = {
        "audit_sha": _audit_sha(vfs),
        "requires_service_active": True,
        "requires_logrotate": False,
    }
    return vfs, targets


def build_hard() -> Tuple[VFS, Dict[str, object]]:
    vfs = _base_fs()
    # Runaway process writing 100 MiB per step to /var/log/app/runaway.log
    vfs.files["/var/log/app/runaway.log"] = VFSFile(
        content=b"", size=int(9.0 * 1024 ** 3)
    )
    vfs.runaway_path = "/var/log/app/runaway.log"
    vfs.runaway_rate = 100 * 1024 * 1024  # 100 MiB/step
    vfs.services["app"] = "failed"
    targets = {
        "audit_sha": _audit_sha(vfs),
        "requires_service_active": True,
        "requires_logrotate": True,
        "logrotate_path": "/etc/logrotate.d/app",
    }
    return vfs, targets


def _audit_sha(vfs: VFS) -> str:
    h = hashlib.sha256()
    for p in sorted(vfs.files):
        if p.startswith("/var/log/audit/"):
            h.update(p.encode())
            h.update(b"\0")
            h.update(vfs.files[p].content)
    return h.hexdigest()


SCENARIOS = {
    "easy": build_easy,
    "medium": build_medium,
    "hard": build_hard,
}

TASK_ORDER = ["easy", "medium", "hard"]