beaver_project/scripts/cleanup-test-users.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path
from typing import Any
from urllib import error as urllib_error
from urllib import parse as urllib_parse
from urllib import request as urllib_request


ROOT_DIR = Path(__file__).resolve().parents[1]
DEFAULT_REGISTRY = ROOT_DIR / "app-instance" / "runtime" / "registry" / "instances.json"
SAFE_PREFIXES = ("smoke", "test", "debug")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Dry-run or purge Beaver smoke/test users.")
    parser.add_argument("--registry", default=str(DEFAULT_REGISTRY), help="Path to app-instance registry JSON.")
    parser.add_argument("--deploy-control-url", default=os.getenv("DEPLOY_CONTROL_URL", "http://127.0.0.1:8090"))
    parser.add_argument("--token", default=os.getenv("DEPLOY_CONTROL_API_TOKEN", ""))
    parser.add_argument("--backend-id", action="append", default=[], help="Explicit backend id to clean.")
    parser.add_argument("--instance-id", action="append", default=[], help="Explicit instance id to clean.")
    parser.add_argument("--username-prefix", action="append", default=[], help="Safe test username prefix, e.g. smoke.")
    parser.add_argument("--execute", action="store_true", help="Actually delete matched instances.")
    parser.add_argument("--purge-data", action="store_true", help="Delete local instance data when executing.")
    parser.add_argument(
        "--keep-user-files",
        action="store_true",
        help="Do not request MinIO/user-file purge when executing.",
    )
    parser.add_argument(
        "--allow-any-prefix",
        action="store_true",
        help="Allow non-test username prefixes. Use only for controlled maintenance.",
    )
    return parser.parse_args()


def load_registry(path: Path) -> list[dict[str, Any]]:
    if not path.exists():
        return []
    payload = json.loads(path.read_text(encoding="utf-8"))
    instances = payload.get("instances", []) if isinstance(payload, dict) else []
    return [item for item in instances if isinstance(item, dict)]


def selected_instances(args: argparse.Namespace) -> list[dict[str, Any]]:
    instances = load_registry(Path(args.registry).expanduser())
    explicit_backend_ids = {item.strip() for item in args.backend_id if item.strip()}
    explicit_instance_ids = {item.strip() for item in args.instance_id if item.strip()}
    prefixes = tuple(item.strip() for item in args.username_prefix if item.strip())

    if not explicit_backend_ids and not explicit_instance_ids and not prefixes:
        raise SystemExit("Refusing cleanup: provide --backend-id, --instance-id, or --username-prefix.")
    unsafe = [prefix for prefix in prefixes if not prefix.startswith(SAFE_PREFIXES)]
    if unsafe and not args.allow_any_prefix:
        raise SystemExit(
            "Refusing cleanup: username prefixes must start with "
            f"{', '.join(SAFE_PREFIXES)} unless --allow-any-prefix is set."
        )

    matches: list[dict[str, Any]] = []
    for item in instances:
        instance_id = str(item.get("instance_id", "") or "")
        backend_id = str(item.get("backend_id", "") or "")
        username = str(item.get("username", "") or "")
        if instance_id in explicit_instance_ids or backend_id in explicit_backend_ids:
            matches.append(item)
            continue
        if prefixes and any(username.startswith(prefix) or backend_id.startswith(prefix) for prefix in prefixes):
            matches.append(item)
    return matches


def delete_instance(args: argparse.Namespace, instance_id: str) -> dict[str, Any]:
    base_url = args.deploy_control_url.rstrip("/")
    url = f"{base_url}/api/instances/{urllib_parse.quote(instance_id, safe='')}"
    headers = {"Accept": "application/json"}
    if args.token.strip():
        headers["Authorization"] = f"Bearer {args.token.strip()}"
    if args.purge_data:
        headers["X-Purge-Data"] = "1"
    if not args.keep_user_files:
        headers["X-Purge-User-Files"] = "1"

    request = urllib_request.Request(url, method="DELETE", headers=headers)
    try:
        with urllib_request.urlopen(request, timeout=120) as response:
            raw = response.read().decode("utf-8")
    except urllib_error.HTTPError as exc:
        detail = exc.reason
        try:
            payload = json.loads(exc.read().decode("utf-8"))
            if isinstance(payload, dict):
                detail = str(payload.get("detail") or detail)
        except Exception:
            pass
        return {"ok": False, "instance_id": instance_id, "error": detail, "status_code": exc.code}
    except urllib_error.URLError as exc:
        return {"ok": False, "instance_id": instance_id, "error": str(exc)}
    if not raw.strip():
        return {"ok": True, "instance_id": instance_id}
    try:
        payload = json.loads(raw)
    except json.JSONDecodeError:
        return {"ok": False, "instance_id": instance_id, "error": "deploy-control response was not valid JSON"}
    return payload if isinstance(payload, dict) else {"ok": False, "instance_id": instance_id, "error": "unexpected response"}


def main() -> int:
    args = parse_args()
    matches = selected_instances(args)
    planned = [
        {
            "instance_id": str(item.get("instance_id", "") or ""),
            "backend_id": str(item.get("backend_id", "") or ""),
            "username": str(item.get("username", "") or ""),
            "instance_root": str(item.get("instance_root", "") or ""),
        }
        for item in matches
    ]
    if not args.execute:
        print(json.dumps({"dry_run": True, "count": len(planned), "planned": planned}, indent=2, ensure_ascii=False))
        return 0

    results = []
    for item in planned:
        instance_id = item["instance_id"]
        if not instance_id:
            results.append({"ok": False, "error": "matched registry record is missing instance_id", "record": item})
            continue
        results.append(delete_instance(args, instance_id))
    ok = all(bool(item.get("ok")) for item in results)
    print(json.dumps({"dry_run": False, "count": len(results), "ok": ok, "results": results}, indent=2, ensure_ascii=False))
    return 0 if ok else 1


if __name__ == "__main__":
    raise SystemExit(main())