safetykit.py

"""
safetykit is a mini-library to demonstrate safety features designed for use in human-in-the-loop scripts.

Running scripts in production without safetykit features is like driving a car without a seatbelt.

Fasten your seatbelt, put on your helmet.
"""

import argparse
import hashlib
import http.server
import json
import os
import secrets
import signal
import socket
import sys
import tempfile
import threading
import time
import urllib.error
import urllib.request
from typing import NoReturn


def praise_dryrun(dryrun: bool = True) -> None:
    """
    The first technique is the dry run mode: default dry run, explicit --wet-run

    Dry run allows an engineer to see what they would do before they actually do it.
    In other words, it **separates planning from execution**.

    The separation of planning from execution is a generally helpful in semi-automated systems.

    **Example**

    ```python
    def gather(paths):
        files = []
        for pattern in paths:
            files.extend(glob.glob(pattern))
        return files

    def execute(files):
        for f in files:
            os.remove(f)

    files = gather([os.path.join(tmp_dir, "*.txt")])
    if dryrun:
        print(f"Would remove: {files}")
    else:
        files = gather([os.path.join(tmp_dir, "*.txt")])
        execute(files)
    ```

    **Claim:** I have not met a script which couldn't trivially be made safer by adding a dry run mode.

    read: https://www.gresearch.com/news/in-praise-of-dry-run
    """

    file_to_remove = os.path.join("/tmp", "backup-2026-05-16.txt")
    demo_content = b"safetykit dryrun demo backup file\n"
    if not os.path.exists(file_to_remove):
        with open(file_to_remove, "wb") as f:
            f.write(demo_content)
    else:
        with open(file_to_remove, "rb") as f:
            existing_content = f.read()
        if existing_content != demo_content:
            raise RuntimeError(f"Refusing to touch existing non-demo file: {file_to_remove}")

    if dryrun:
        print(f"Would remove: {file_to_remove}")
    else:
        print("Wet run enabled. Removing target now.")
        os.remove(file_to_remove)
        print(f"Removed: {file_to_remove}")


def ask_for_confirmation() -> bool:
    """
    The second technique is asking for confirmation before a destructive action.
    You've used this hundreds of times on SaaS and CSP dashboards. Confirmation
    boxes exists because they work!

    Confirmation boxes help users avoid what's called a 'slip'. A slip is when a user
    accidentally performs an action they didn't intend to. Humans attention is flaky, especially
    these days, and the confirmation software feature significantly reduces regrettable actions.

    It's IMPORTANT to know that Y/N confirmation is significantly less effective than 'type this name' confirmation.
    If the user is doing something destructive and irreversible, Y/N confirmation is not effective enough.
    Encouraging mindless Y,Y,Y or tab-tab-tab tapping is a receipe for disaster.

    The Therac-25 accident involved software which trained operators to mindlessly press the P key (Proceed : continue).
    And they pressed it even when the machine was indirectly asking if they wanted to kill the patient.

    read: Nancy G. Leveson and Clark S. Turner, "An Investigation of the Therac-25 Accidents,"
          IEEE Computer, Vol. 26, No. 7, July 1993, pp. 18-41.
    """
    import os
    import random

    tmp_dir = "/tmp"
    files = [f for f in os.listdir(tmp_dir) if os.path.isfile(os.path.join(tmp_dir, f))]
    if not files:
        print("No files in /tmp to delete.")
        return False

    file_to_remove = random.choice(files)
    print(f"Selected file for deletion: {file_to_remove}")
    confirmation = input(f"Type the filename '{file_to_remove}' to confirm deletion: ")
    if confirmation == file_to_remove:
        os.remove(os.path.join(tmp_dir, file_to_remove))
        print(f"✅ Removed: {file_to_remove}")
        return True
    else:
        print("❌ Confirmation failed. File not removed.")
        return False


def pause_a_beat():
    """
    The third technique involves slowing down execution to give users a chance to review the action
    or rethink their intentions.

    You've surely in your life panic-mashed CTRL-C to stop a computer doing something you just realized
    you didn't want. If the computer was operating at full gallop, your 100+ms monkey reflexes were probably too slow.

    By adding pauses into execution, a human operator is able to:
    1. Re-evaluate intentions *after* initiation.
    2. Keep up with the computer's pace and intervene.

    This technique is best suited when a script's wall clock performance is not at all important. For example,
    if a dangerous script is run once a week and takes 10 seconds to run with pauses included, the gain in safety
    is worth the 10 seconds of extra operator time.
    """

    print("🌊 Running rm -rf on system root (/).\nLet's free some disk space!")
    try:
        time.sleep(2)
    except KeyboardInterrupt:
        print("glad I paused?")
        return
    print("🌵 Just kidding!")


def abort_safely():
    """
    The fourth technique is a demonstration of a classic safety technique, used in factories
    and airplanes all over the world: abort or fail into a safe state. "fail-safe" design.

    A most famous example is Otis's elevator: emergency brakes in elevators use springs that clamp
    onto guide rails when cable tension is lost.

    In the domain of scripts, it's more relevant to talk about abort-safety.
    Abort-safety is the presence of safety when a script is cancelled (SIGINT), terminated (SIGTERM), or killed (SIGKILL).

    **Example:** Safely writing to a file such that ctrl-C or kill won't corrupt the file.
    This function uses a temp file + an atomic move, and register handlers for SIGINT/SIGTERM to clean up.

    read: https://en.wikipedia.org/wiki/Fail-safe
    """
    data_to_write = "Important data that shouldn't be partially saved!\n"
    output_path = os.path.join(tempfile.gettempdir(), "critical_output.txt")
    temp_path = output_path + ".tmp"

    def cleanup_temp(*_):
        try:
            os.remove(temp_path)
        except OSError:
            pass
        sys.exit(1)  # exit unhappily, but safely

    # Register abort/termination signal handlers
    for sig in (signal.SIGINT, signal.SIGTERM):
        signal.signal(sig, cleanup_temp)

    print(f"Writing data to {output_path} safely (interrupt with Ctrl-C or kill)...")
    try:
        with open(temp_path, "w") as f:
            f.write(data_to_write)
            f.flush()
            os.fsync(f.fileno())
            # Intentionally sleep to give user a chance to Ctrl+C
            print("...writing (simulate slow write, try interrupting now)...")
            time.sleep(5)
        # Perform atomic move only if everything completed
        os.replace(temp_path, output_path)
        print(f"✅ Safely wrote data to {output_path}")
    except Exception as exc:
        print(f"❌ Exception during write: {exc}")
    finally:
        cleanup_temp()


def undo() -> bool:
    """
    The fifth technique is making destructive operations undoable.

    Instead of deleting a file immediately, move it into quarantine first. Give
    the human a short cancellation window, then commit the delete only after
    that window closes.
    """
    workspace = tempfile.mkdtemp(prefix="safetykit_undo_", dir="/tmp")
    quarantine_dir = os.path.join(workspace, "quarantine")
    os.mkdir(quarantine_dir)

    original_path = os.path.join(workspace, "important-report.txt")
    quarantined_path = os.path.join(quarantine_dir, "important-report.txt.deleted")

    with open(original_path, "w") as f:
        f.write("important report contents\n")

    print(f"Demo workspace: {workspace}")
    print(f"Created: {original_path}")

    os.replace(original_path, quarantined_path)
    print(f"Moved to quarantine: {quarantined_path}")
    print("Press any key within 10 seconds to undo the delete.")

    def read_key_with_timeout(seconds: int) -> str | None:
        if not sys.stdin.isatty():
            time.sleep(seconds)
            return None

        import select
        import termios
        import tty

        fd = sys.stdin.fileno()
        old_settings = termios.tcgetattr(fd)
        try:
            tty.setcbreak(fd)
            ready, _, _ = select.select([sys.stdin], [], [], seconds)
            if ready:
                return sys.stdin.read(1)
            return None
        finally:
            termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)

    for remaining in range(10, 0, -1):
        print(f"Committing delete in {remaining:2d}s...")
        key = read_key_with_timeout(1)
        if key is not None:
            print()
            os.replace(quarantined_path, original_path)
            print(f"✅ Cancelled delete and restored: {original_path}")
            return False

    os.remove(quarantined_path)
    print(f"Committed delete: {original_path}")
    return True


def feedback() -> None:
    """
    The sixth technique is visible feedback while a script is doing something risky.

    This demo ramps synthetic CPU pressure and prints colour-coded status lines.
    The point is not benchmarking; it is making the system's changing state hard
    to miss while still stopping cleanly on Ctrl-C.
    """
    core_count = os.cpu_count() or 1
    selected_core = min(7, core_count - 1)
    targets = [
        (10, "blue", "94"),
        (30, "green", "92"),
        (50, "yellow", "93"),
        (70, "orange", "38;5;208"),
        (90, "red", "91"),
        (100, "red", "91"),
    ]

    print(f"Selected logical core {selected_core} from {core_count} available cores.")
    print("Raising CPU pressure. Press Ctrl-C to stop.\n")

    if hasattr(os, "sched_setaffinity"):
        try:
            os.sched_setaffinity(0, {selected_core})
            print(f"Pinned process to logical core {selected_core}.")
        except OSError:
            print("CPU affinity pinning is not available on this platform; running best-effort.")
    else:
        print("CPU affinity pinning is not available on this platform; running best-effort.")

    def pressure_cycle(target_percent: int, duration: float = 0.45) -> None:
        busy_for = duration * target_percent / 100
        idle_for = max(0, duration - busy_for)
        start = time.perf_counter()
        while time.perf_counter() - start < busy_for:
            pass
        if idle_for:
            time.sleep(idle_for)

    try:
        while True:
            for target, label, colour in targets:
                pressure_cycle(target)
                usage = min(99.9, max(0.0, target - 0.3 + (target / 137)))
                filled = min(28, round(target / 100 * 28))
                bar = "#" * filled + "-" * (28 - filled)
                print(
                    f"\033[{colour}m"
                    f"core={selected_core:02d} target={target:3d}% usage={usage:5.1f}% {label:<6} [{bar}]"
                    f"\033[0m"
                )
    except KeyboardInterrupt:
        print("\nRecieved interrupt. Stopping...")


def audit() -> None:
    """
    The seventh technique is leaving durable evidence behind.

    This demo writes JSON Lines audit events and hash-chains them so each event
    commits to the previous event's hash.
    """
    audit_path = os.path.join("/tmp", "safetykit-audit.jsonl")
    run_id = secrets.token_hex(8)
    previous_hash = "0" * 64
    actions = [
        ("plan", "created"),
        ("review", "approved"),
        ("execute", "skipped"),
    ]

    print(f"Writing audit events to: {audit_path}")
    with open(audit_path, "a") as f:
        for sequence, (action, outcome) in enumerate(actions, start=1):
            event = {
                "run_id": run_id,
                "sequence": sequence,
                "action": action,
                "outcome": outcome,
                "previous_hash": previous_hash,
                "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            }
            event_json = json.dumps(event, sort_keys=True, separators=(",", ":"))
            event_hash = hashlib.sha256(event_json.encode("utf-8")).hexdigest()
            event["event_hash"] = event_hash
            f.write(json.dumps(event, sort_keys=True, separators=(",", ":")) + "\n")
            previous_hash = event_hash

    print(f"Audit run recorded with run_id={run_id}")


class TwoPersonRuleError(Exception):
    pass


def two_person_rule(
    secret: str | None = None,
    initiator_url: str | None = None,
    timeout_seconds: int = 300,
) -> bool:
    """
    Per US Air Force Instruction (AFI) 91-104, "the two-person concept" is designed
    to prevent accidental or malicious launch of nuclear weapons by a single individual.

    We can do this with software scripts too!

    This function implements a coordination mechanism where two people must independently
    run the script and confirm each other before one script proceeds.

    Args:
        secret: A shared secret known to both persons. If None, reads from
                TWO_PERSON_SECRET env var or prompts for input.
        initiator_url: If provided, run as responder connecting to this URL.
                       If None, run as initiator (start server).
        timeout_seconds: How long to wait for the second person (default 5 minutes).

    Returns:
        True if this script should proceed (initiator after confirmation).
        False and exits if this script is the responder (other script proceeds).

    Raises:
        TwoPersonRuleError: If coordination fails (timeout, wrong secret, etc.)

    ref: https://en.wikipedia.org/wiki/Two-person_rule
    """

    def get_local_ip() -> str:
        """Get the local IP address that can be reached by other machines."""
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
                s.connect(("8.8.8.8", 80))
                return s.getsockname()[0]
        except Exception:
            return "127.0.0.1"

    def run_as_initiator(secret_hash: str) -> bool:
        """Run as the initiator: start HTTP server and wait for responder."""
        state = {
            "confirmed": False,
            "error": None,
            "responder_nonce": None,
            "initiator_nonce": secrets.token_hex(16),
            "ack_received": False,
        }
        state_lock = threading.Lock()
        confirmation_event = threading.Event()

        class TwoPersonHandler(http.server.BaseHTTPRequestHandler):
            def log_message(self, format, *args):
                pass

            def do_POST(self):
                if self.path == "/join":
                    self._handle_join()
                elif self.path == "/ack":
                    self._handle_ack()
                else:
                    self.send_error(404)

            def _handle_join(self):
                try:
                    content_length = int(self.headers.get("Content-Length", 0))
                    body = self.rfile.read(content_length).decode("utf-8")
                    data = json.loads(body)

                    responder_secret_hash = data.get("secret_hash", "")
                    responder_nonce = data.get("nonce", "")

                    if not secrets.compare_digest(responder_secret_hash, secret_hash):
                        self.send_response(403)
                        self.send_header("Content-Type", "application/json")
                        self.end_headers()
                        self.wfile.write(json.dumps({"error": "Secret mismatch"}).encode())
                        with state_lock:
                            state["error"] = "Responder provided incorrect secret"
                        return

                    with state_lock:
                        state["responder_nonce"] = responder_nonce
                        state["confirmed"] = True

                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.end_headers()
                    response = {
                        "status": "confirmed",
                        "initiator_nonce": state["initiator_nonce"],
                    }
                    self.wfile.write(json.dumps(response).encode())

                except Exception as e:
                    self.send_error(400, str(e))

            def _handle_ack(self):
                try:
                    content_length = int(self.headers.get("Content-Length", 0))
                    body = self.rfile.read(content_length).decode("utf-8")
                    data = json.loads(body)

                    ack_nonce = data.get("initiator_nonce", "")
                    if secrets.compare_digest(ack_nonce, state["initiator_nonce"]):
                        with state_lock:
                            state["ack_received"] = True
                        confirmation_event.set()

                    self.send_response(200)
                    self.send_header("Content-Type", "application/json")
                    self.end_headers()
                    self.wfile.write(json.dumps({"status": "ok"}).encode())

                except Exception as e:
                    self.send_error(400, str(e))

        server = http.server.HTTPServer(("0.0.0.0", 0), TwoPersonHandler)
        port = server.server_address[1]
        local_ip = get_local_ip()

        print("=" * 60)
        print("TWO-PERSON RULE: Initiator Mode")
        print("=" * 60 + "\n")
        print("Share the following with the second person:")
        print(f"  URL: http://{local_ip}:{port}")
        print("  (They must also know the shared secret)\n")
        print(f"Waiting for second person to connect (timeout: {timeout_seconds}s)...")
        print("Press Ctrl+C to abort.\n")

        server_thread = threading.Thread(target=server.serve_forever)
        server_thread.daemon = True
        server_thread.start()

        try:
            if not confirmation_event.wait(timeout=timeout_seconds):
                server.shutdown()
                raise TwoPersonRuleError(f"Timeout: No second person connected within {timeout_seconds} seconds.")

            server.shutdown()

            with state_lock:
                if state["error"]:
                    raise TwoPersonRuleError(state["error"])

                if not state["ack_received"]:
                    raise TwoPersonRuleError("Handshake incomplete: no acknowledgment received.")

            print("\n" + ("=" * 60))
            print("Two-person confirmation COMPLETE!")
            print("Second person has confirmed. This script will proceed.")
            print(("=" * 60) + "\n")
            return True

        except KeyboardInterrupt:
            server.shutdown()
            raise TwoPersonRuleError("Aborted by user (Ctrl+C).")

    def run_as_responder(secret_hash: str) -> NoReturn:
        """Run as the responder: connect to initiator and confirm."""
        print("=" * 60)
        print("TWO-PERSON RULE: Responder Mode")
        print("=" * 60)
        print(f"\nConnecting to initiator at: {initiator_url}\n")

        responder_nonce = secrets.token_hex(16)

        join_url = initiator_url.rstrip("/") + "/join"
        join_data = json.dumps(
            {
                "secret_hash": secret_hash,
                "nonce": responder_nonce,
            }
        ).encode("utf-8")

        try:
            req = urllib.request.Request(
                join_url,
                data=join_data,
                headers={"Content-Type": "application/json"},
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=30) as response:
                response_data = json.loads(response.read().decode("utf-8"))

        except urllib.error.HTTPError as e:
            if e.code == 403:
                raise TwoPersonRuleError("Secret mismatch! Your secret does not match the initiator's secret.")
            raise TwoPersonRuleError(f"Connection failed: HTTP {e.code}")
        except urllib.error.URLError as e:
            raise TwoPersonRuleError(f"Connection failed: {e.reason}")
        except Exception as e:
            raise TwoPersonRuleError(f"Connection failed: {e}")

        if response_data.get("status") != "confirmed":
            raise TwoPersonRuleError(f"Unexpected response: {response_data}")

        initiator_nonce = response_data.get("initiator_nonce", "")
        if not initiator_nonce:
            raise TwoPersonRuleError("Invalid response: missing initiator nonce")

        print("Secret verified by initiator!")
        print()

        ack_url = initiator_url.rstrip("/") + "/ack"
        ack_data = json.dumps(
            {
                "initiator_nonce": initiator_nonce,
            }
        ).encode("utf-8")

        try:
            req = urllib.request.Request(
                ack_url,
                data=ack_data,
                headers={"Content-Type": "application/json"},
                method="POST",
            )
            with urllib.request.urlopen(req, timeout=30) as response:
                response.read()
        except Exception as e:
            raise TwoPersonRuleError(f"Acknowledgment failed: {e}")

        print("=" * 60)
        print("Two-person confirmation COMPLETE!")
        print("Initiator has been notified. THIS script will now EXIT.")
        print("(The initiator's script will proceed.)")
        print("=" * 60)
        print()
        sys.exit(0)

    # Get the shared secret
    if secret is None:
        secret = os.environ.get("TWO_PERSON_SECRET")
    if secret is None:
        print("Two-Person Rule: Enter the shared secret known to both persons.")
        secret = input("Secret: ").strip()
    if not secret:
        raise TwoPersonRuleError("A shared secret is required for two-person coordination.")

    # Hash the secret to avoid transmitting it in plaintext
    secret_hash = hashlib.sha256(secret.encode()).hexdigest()

    if initiator_url is None:
        confirmed = run_as_initiator(secret_hash)
    else:
        run_as_responder(secret_hash)

    if confirmed:
        print("Demo: The two-person rule was satisfied!")
        print("In a real script, dangerous operations would proceed here.")
    return confirmed


def main():
    parser = argparse.ArgumentParser(description="SafetyKit Demos: Choose which safety technique to demo.")
    parser.add_argument(
        "demo",
        choices=["dryrun", "confirm", "pause", "abort", "undo", "feedback", "audit", "two_person"],
        help=(
            "Which demo to run: 'dryrun' (dry run mode), 'confirm' (ask for confirmation), "
            "'pause' (pause before action), 'abort' (abort-safe writes), 'undo' (undo window), "
            "'feedback' (visible progress and interrupt handling), 'audit' (hash-chained audit log), "
            "'two_person' (two-person rule coordination)."
        ),
    )
    parser.add_argument(
        "--connect",
        metavar="URL",
        help="For 'two_person' demo: connect to an initiator at this URL (responder mode).",
    )
    parser.add_argument(
        "--secret",
        help="For 'two_person' demo: the shared secret (can also use TWO_PERSON_SECRET env var).",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=300,
        help="For 'two_person' demo: timeout in seconds waiting for second person (default: 300).",
    )
    parser.add_argument(
        "--wet-run",
        action="store_true",
        help="For 'dryrun' demo: actually remove the demo file instead of only previewing it.",
    )

    args = parser.parse_args()
    if args.demo == "dryrun":
        praise_dryrun(dryrun=not args.wet_run)
    elif args.demo == "confirm":
        ask_for_confirmation()
    elif args.demo == "pause":
        pause_a_beat()
    elif args.demo == "abort":
        abort_safely()
    elif args.demo == "undo":
        undo()
    elif args.demo == "feedback":
        feedback()
    elif args.demo == "audit":
        audit()
    elif args.demo == "two_person":
        try:
            two_person_rule(
                secret=args.secret,
                initiator_url=args.connect,
                timeout_seconds=args.timeout,
            )
        except TwoPersonRuleError as e:
            print(f"Two-person rule failed: {e}")
            sys.exit(1)


if __name__ == "__main__":
    main()