kill_switch_worker.py

#!/usr/bin/env python3
"""
machineid-runtime-kill-switch
Terminal 2: operator "kill switch" worker.

Purpose:
- Run this in a second terminal while runaway_agent.py is running.
- Revoke / unrevoke a device on demand (runtime kill switch).
- Optionally confirm state via POST /validate (canonical).

Env vars:
- MACHINEID_ORG_KEY        (required) org_...
- MACHINEID_DEVICE_ID      (optional) default: kill-switch-agent-01
- MACHINEID_BASE_URL       (optional) default: https://machineid.io
"""

import os
import sys
import json
from datetime import datetime, timezone

import requests


def env(name: str, default: str | None = None) -> str | None:
    v = os.getenv(name)
    if v is None:
        return default
    v = v.strip()
    return v if v else default


def must_env(name: str) -> str:
    v = env(name)
    if not v:
        print(f"[fatal] Missing required env var: {name}", file=sys.stderr)
        sys.exit(2)
    return v


def iso_now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def post_json(url: str, headers: dict, payload: dict, timeout_s: int = 12) -> dict:
    r = requests.post(url, headers=headers, json=payload, timeout=timeout_s)
    try:
        data = r.json()
    except Exception:
        data = {"status": "error", "error": f"Non-JSON response (HTTP {r.status_code})"}
    if r.status_code >= 400:
        if isinstance(data, dict) and data.get("error"):
            return {"status": "error", "error": data.get("error"), "http": r.status_code}
        return {"status": "error", "error": f"HTTP {r.status_code}", "http": r.status_code, "body": data}
    return data


def print_validate(val: dict) -> None:
    allowed = val.get("allowed")
    code = val.get("code")
    request_id = val.get("request_id")
    plan = val.get("planTier")
    used = val.get("devicesUsed")
    limit = val.get("limit")
    suffix = ""
    if plan is not None or (used is not None and limit is not None):
        suffix = f" | plan={plan} used={used}/{limit}"
    print(f"[{iso_now()}] validate: allowed={allowed} code={code} request_id={request_id}{suffix}")


def main() -> None:
    org_key = must_env("MACHINEID_ORG_KEY")
    device_id = env("MACHINEID_DEVICE_ID", "kill-switch-agent-01")
    base_url = env("MACHINEID_BASE_URL", "https://machineid.io").rstrip("/")

    if not org_key.startswith("org_"):
        print("[fatal] MACHINEID_ORG_KEY must start with org_", file=sys.stderr)
        sys.exit(2)

    headers = {"Content-Type": "application/json", "x-org-key": org_key}

    revoke_url = f"{base_url}/api/v1/devices/revoke"
    unrevoke_url = f"{base_url}/api/v1/devices/unrevoke"
    validate_url = f"{base_url}/api/v1/devices/validate"

    print(f"[{iso_now()}] kill switch worker ready")
    print(f"[{iso_now()}] base_url={base_url}")
    print(f"[{iso_now()}] device_id={device_id}")
    print("")
    print("Commands:")
    print("  kill     -> POST /devices/revoke")
    print("  restore  -> POST /devices/unrevoke")
    print("  status   -> POST /devices/validate")
    print("  help     -> show commands")
    print("  quit     -> exit")
    print("")

    while True:
        try:
            cmd = input("> ").strip().lower()
        except (EOFError, KeyboardInterrupt):
            print("\n[exit] bye")
            return

        if cmd in ("q", "quit", "exit"):
            print("[exit] bye")
            return

        if cmd in ("h", "help", "?"):
            print("kill | restore | status | quit")
            continue

        if cmd in ("status", "s"):
            val = post_json(validate_url, headers, {"deviceId": device_id})
            print_validate(val)
            continue

        if cmd in ("kill", "k", "revoke"):
            out = post_json(revoke_url, headers, {"deviceId": device_id})
            print(f"[{iso_now()}] revoke response:")
            print(json.dumps(out, indent=2, sort_keys=True))
            val = post_json(validate_url, headers, {"deviceId": device_id})
            print_validate(val)
            continue

        if cmd in ("restore", "r", "unrevoke"):
            out = post_json(unrevoke_url, headers, {"deviceId": device_id})
            print(f"[{iso_now()}] unrevoke response:")
            print(json.dumps(out, indent=2, sort_keys=True))
            val = post_json(validate_url, headers, {"deviceId": device_id})
            print_validate(val)
            continue

        print("Unknown command. Type 'help'.")


if __name__ == "__main__":
    main()