cang-ying/remote-node/tools/cangerctl.py

#!/usr/bin/env python3
"""Small CLI used by the persona and owner on the Singapore server."""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from typing import Any
from urllib.error import HTTPError
from urllib.request import Request, urlopen


def request(
    base: str, token: str, method: str, path: str, body: dict[str, Any] | None = None
) -> dict[str, Any]:
    data = None if body is None else json.dumps(body).encode("utf-8")
    headers = {"Authorization": f"Bearer {token}", "Accept": "application/json"}
    if data is not None:
        headers["Content-Type"] = "application/json"
    req = Request(base.rstrip("/") + path, data=data, headers=headers, method=method)
    try:
        with urlopen(req, timeout=30) as response:
            return json.load(response)
    except HTTPError as exc:
        print(exc.read().decode("utf-8", errors="replace"), file=sys.stderr)
        raise


def token_for(role: str) -> str:
    name = {
        "admin": "CANGER_ADMIN_TOKEN",
        "owner": "CANGER_OWNER_TOKEN",
        "persona": "CANGER_PERSONA_TOKEN",
        "operator": "CANGER_OPERATOR_TOKEN",
    }[role]
    value = os.environ.get(name, "")
    if not value:
        raise RuntimeError(f"{name} is not set")
    return value


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--url", default=os.environ.get("CANGER_CONTROL_URL", "http://127.0.0.1:8787")
    )
    sub = parser.add_subparsers(dest="command", required=True)
    sub.add_parser("pair-code")
    sub.add_parser("nodes")
    task = sub.add_parser("task")
    task.add_argument("--node", required=True)
    task.add_argument("--action", required=True)
    task.add_argument("--args", default="{}")
    approve = sub.add_parser("approve")
    approve.add_argument("task_id")
    reject = sub.add_parser("reject")
    reject.add_argument("task_id")
    show = sub.add_parser("show")
    show.add_argument("task_id")
    watch = sub.add_parser("watch")
    watch.add_argument("task_id")
    grant = sub.add_parser("grant")
    grant.add_argument("--node", required=True)
    grant.add_argument("--hours", type=int, default=4)
    grant.add_argument("--reason", default="苍耳人格体编程、调试与诊断")
    sub.add_parser("grants")
    revoke = sub.add_parser("revoke")
    revoke.add_argument("grant_id")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.command == "pair-code":
        result = request(args.url, token_for("admin"), "POST", "/v1/pairings", {})
    elif args.command == "nodes":
        result = request(args.url, token_for("operator"), "GET", "/v1/nodes")
    elif args.command == "task":
        result = request(
            args.url,
            token_for("operator"),
            "POST",
            "/v1/tasks",
            {
                "node_id": args.node,
                "action": args.action,
                "args": json.loads(args.args),
            },
        )
    elif args.command == "approve":
        result = request(
            args.url,
            token_for("owner"),
            "POST",
            f"/v1/tasks/{args.task_id}/approve",
            {},
        )
    elif args.command == "reject":
        result = request(
            args.url,
            token_for("owner"),
            "POST",
            f"/v1/tasks/{args.task_id}/reject",
            {},
        )
    elif args.command == "show":
        result = request(
            args.url,
            token_for("operator"),
            "GET",
            f"/v1/tasks/{args.task_id}",
        )
    elif args.command == "watch":
        after = 0
        while True:
            events = request(
                args.url,
                token_for("operator"),
                "GET",
                f"/v1/tasks/{args.task_id}/events?after={after}",
            )
            for event in events["events"]:
                after = max(after, event["seq"])
                print(
                    f"[{event['seq']}] {event['kind']}: "
                    + json.dumps(event["payload"], ensure_ascii=False),
                    flush=True,
                )
            state = request(
                args.url,
                token_for("operator"),
                "GET",
                f"/v1/tasks/{args.task_id}",
            )
            if state["status"] in {"completed", "rejected"}:
                result = state
                break
            time.sleep(1)
    elif args.command == "grant":
        result = request(
            args.url,
            token_for("operator"),
            "POST",
            "/v1/grants",
            {
                "node_id": args.node,
                "duration_seconds": args.hours * 3600,
                "reason": args.reason,
            },
        )
    elif args.command == "grants":
        result = request(args.url, token_for("operator"), "GET", "/v1/grants")
    else:
        result = request(
            args.url,
            token_for("owner"),
            "POST",
            f"/v1/grants/{args.grant_id}/revoke",
            {},
        )
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()