112 lines
3.5 KiB
Python

"""Shared protocol rules for the Cang'er remote node."""
from __future__ import annotations
import hashlib
import hmac
import json
import re
import secrets
import time
from typing import Any
MAX_JSON_BYTES = 256 * 1024
MAX_RESULT_BYTES = 512 * 1024
PAIRING_TTL_SECONDS = 10 * 60
TASK_LEASE_SECONDS = 5 * 60
ACTION_NAMES = {
"system.status",
"repo.status",
"repo.fetch",
"service.logs",
"command.run",
}
NODE_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$")
SERVICE_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.@-]{0,127}$")
def now_ts() -> int:
return int(time.time())
def new_id(prefix: str) -> str:
return f"{prefix}_{secrets.token_hex(12)}"
def new_token() -> str:
return secrets.token_urlsafe(32)
def new_pairing_code() -> str:
alphabet = "ABCDEFGHJKLMNPQRSTUVWXYZ23456789"
return "-".join(
"".join(secrets.choice(alphabet) for _ in range(4)) for _ in range(3)
)
def secret_hash(value: str) -> str:
return hashlib.sha256(value.encode("utf-8")).hexdigest()
def secret_matches(value: str, expected_hash: str) -> bool:
return hmac.compare_digest(secret_hash(value), expected_hash)
def compact_json(value: Any) -> str:
return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True)
def validate_node_name(value: Any) -> str:
if not isinstance(value, str) or not NODE_NAME_RE.fullmatch(value):
raise ValueError("node_name must use 1-64 letters, digits, dot, underscore or dash")
return value
def validate_task(action: Any, args: Any) -> tuple[str, dict[str, Any]]:
if not isinstance(action, str) or action not in ACTION_NAMES:
raise ValueError(f"unsupported action: {action!r}")
if args is None:
args = {}
if not isinstance(args, dict):
raise ValueError("args must be an object")
if action == "system.status":
if args:
raise ValueError("system.status does not accept args")
elif action in {"repo.status", "repo.fetch"}:
if set(args) != {"repo"} or not isinstance(args["repo"], str):
raise ValueError(f"{action} requires string arg: repo")
elif action == "service.logs":
if not isinstance(args.get("service"), str) or not SERVICE_NAME_RE.fullmatch(
args["service"]
):
raise ValueError("service.logs requires a safe service name")
lines = args.get("lines", 200)
if not isinstance(lines, int) or not 1 <= lines <= 2000:
raise ValueError("service.logs lines must be between 1 and 2000")
args = {"service": args["service"], "lines": lines}
elif action == "command.run":
argv = args.get("argv")
cwd = args.get("cwd")
timeout = args.get("timeout_seconds", 300)
if (
not isinstance(argv, list)
or not argv
or len(argv) > 64
or any(not isinstance(item, str) or len(item) > 4096 for item in argv)
):
raise ValueError("command.run argv must be a non-empty string array")
if not isinstance(cwd, str):
raise ValueError("command.run requires string arg: cwd")
if not isinstance(timeout, int) or not 1 <= timeout <= 1800:
raise ValueError("command.run timeout_seconds must be between 1 and 1800")
args = {"argv": argv, "cwd": cwd, "timeout_seconds": timeout}
encoded = compact_json(args).encode("utf-8")
if len(encoded) > MAX_JSON_BYTES:
raise ValueError("task args are too large")
return action, args