runaway_agent.py

#!/usr/bin/env python3
"""
machineid-runtime-kill-switch
Runaway agent demo.

Behavior:
- Registers once on startup (idempotent)
- Loops forever:
  - POST /validate
  - if allowed == false -> exit immediately
  - otherwise, print heartbeat and sleep

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

import os
import sys
import time
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:
        return {
            "allowed": False,
            "code": "HTTP_ERROR",
            "request_id": None,
            "http": r.status_code,
            "body": data,
        }

    return data


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("/")
    heartbeat_raw = env("HEARTBEAT_SECONDS", "5")

    try:
        heartbeat_s = max(1, int(heartbeat_raw))
    except Exception:
        heartbeat_s = 5

    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}

    reg_url = f"{base_url}/api/v1/devices/register"
    val_url = f"{base_url}/api/v1/devices/validate"

    print(f"[{iso_now()}] runaway agent starting")
    print(f"[{iso_now()}] base_url={base_url} device_id={device_id}")

    # Register (idempotent)
    reg = post_json(reg_url, headers, {"deviceId": device_id})
    if reg.get("status") not in ("ok", "exists"):
        print(f"[{iso_now()}] [fatal] register failed: {reg}", file=sys.stderr)
        sys.exit(1)

    print(f"[{iso_now()}] register {reg.get('status')}")

    n = 0
    while True:
        n += 1

        val = post_json(val_url, headers, {"deviceId": device_id})

        if val.get("allowed") is False:
            print(
                f"[{iso_now()}] [stop] execution denied: "
                f"{val.get('code')} {val.get('request_id')}"
            )
            sys.exit(0)

        devices_used = val.get("devicesUsed")
        limit = val.get("limit")
        plan = val.get("planTier")

        suffix = ""
        if plan is not None and devices_used is not None and limit is not None:
            suffix = f" | plan={plan} used={devices_used}/{limit}"

        print(f"[{iso_now()}] heartbeat #{n}: allowed=true{suffix}")
        time.sleep(heartbeat_s)


if __name__ == "__main__":
    main()