aios/tools/aios_spark_runner/runnerctl.py

#!/usr/bin/env python3
"""CLI and public API for the managed AIOS isolated Python runner."""
from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
import pwd
import sys
if __package__ in {None, ""}:
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

try:
    from tools.aios_spark_runner.artifact import build_artifact
    from tools.aios_spark_runner.config import RunnerConfig, load_config, validate_artifact
    from tools.aios_spark_runner.errors import ConfigError
    from tools.aios_spark_runner.installer import provision, rollback, uninstall
    from tools.aios_spark_runner.provenance import verify_manifest
    from tools.aios_spark_runner.broker import make_secret_ref_probe
    from tools.aios_spark_runner.runtime import (
        check_heartbeat,
        probe_worktree,
        run_forever,
        utc_now,
        write_heartbeat,
    )
    from tools.aios_spark_runner.securefs import sha256_file
except ModuleNotFoundError:
    from aios_spark_runner.artifact import build_artifact
    from aios_spark_runner.config import RunnerConfig, load_config, validate_artifact
    from aios_spark_runner.errors import ConfigError
    from aios_spark_runner.installer import provision, rollback, uninstall
    from aios_spark_runner.provenance import verify_manifest
    from aios_spark_runner.broker import make_secret_ref_probe
    from aios_spark_runner.runtime import (
        check_heartbeat,
        probe_worktree,
        run_forever,
        utc_now,
        write_heartbeat,
    )
    from aios_spark_runner.securefs import sha256_file


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="aios-spark-runner")
    commands = parser.add_subparsers(dest="command", required=True)

    build = commands.add_parser("build-artifact")
    build.add_argument("--commit", required=True)
    build.add_argument("--output", type=Path, required=True)
    build.add_argument("--signing-key", type=Path, required=True)
    build.add_argument("--source-repo", type=Path, default=Path.cwd())

    install = commands.add_parser("provision")
    install.add_argument("--artifact", type=Path, required=True)
    install.add_argument("--manifest", type=Path, required=True)
    install.add_argument("--trusted-public-key", type=Path, required=True)
    install.add_argument("--source-repo", type=Path, required=True)
    install.add_argument("--config", type=Path, required=True)
    install.add_argument("--root", type=Path, default=Path("/"))
    install.add_argument("--test-mode", action="store_true")

    undo = commands.add_parser("rollback")
    undo.add_argument("--root", type=Path, default=Path("/"))
    undo.add_argument("--receipt", type=Path, required=True)
    undo.add_argument("--test-mode", action="store_true")

    remove = commands.add_parser("uninstall")
    remove.add_argument("--root", type=Path, default=Path("/"))
    remove.add_argument("--test-mode", action="store_true")

    accept = commands.add_parser("acceptance")
    accept.add_argument("--config", type=Path, required=True)
    accept.add_argument("--bare-repo", type=Path, required=True)
    accept.add_argument("--commit", required=True)
    accept.add_argument("--probe-parent", type=Path, required=True)
    accept.add_argument("--heartbeat", type=Path, required=True)
    accept.add_argument("--max-age-seconds", type=int, default=300)

    run = commands.add_parser("run")
    run.add_argument("--config", type=Path, required=True)
    run.add_argument("--interval-seconds", type=int, default=30)
    return parser


def _execute(args: argparse.Namespace) -> object:
    if args.command == "build-artifact":
        return build_artifact(
            args.commit,
            args.output,
            signing_key=args.signing_key,
            source_repo=args.source_repo,
        )
    if args.command == "provision":
        uid = os.getuid() if args.test_mode else 0
        return provision(
            args.artifact,
            args.manifest,
            args.trusted_public_key,
            args.source_repo,
            args.config,
            args.root,
            test_mode=args.test_mode,
            expected_uid=uid,
        )
    if args.command == "rollback":
        return rollback(args.root, args.receipt, test_mode=args.test_mode)
    if args.command == "uninstall":
        return uninstall(args.root, test_mode=args.test_mode)
    if args.command == "acceptance":
        config = load_config(args.config, expected_uid=0)
        heartbeat = check_heartbeat(
            config,
            args.heartbeat,
            max_age_seconds=args.max_age_seconds,
            expected_uid=runner_uid(config.runner_run_user),
        )
        probe = probe_worktree(args.bare_repo, args.commit, args.probe_parent)
        return {"heartbeat": heartbeat["status"], **probe}
    if args.command == "run":
        config = load_config(
            args.config, expected_uid=0, expected_modes=(0o640,)
        )
        return run_forever(config, interval_seconds=args.interval_seconds)
    raise ConfigError("unknown runner command")


def runner_uid(name: str) -> int:
    try:
        return pwd.getpwnam(name).pw_uid
    except KeyError as exc:
        raise ConfigError("runner system user is unavailable") from exc


def main(argv: list[str] | None = None) -> int:
    try:
        result = _execute(_parser().parse_args(argv))
    except ConfigError as exc:
        print(str(exc), file=sys.stderr)
        return 2
    if isinstance(result, int):
        return result
    print(json.dumps(result, sort_keys=True))
    return 0


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