ZKAC/cli/zkac_cli/main.py

"""zkac-node CLI entrypoint."""

from __future__ import annotations

import argparse
import base64
import json
import secrets
import shutil
import sys
from pathlib import Path

import zkac

from zkac_cli import creds
from zkac_cli import client_ops
from zkac_cli import paths
from zkac_cli import registry_local
from zkac_cli import server_app


MGMT_ROLE = server_app.MGMT_ROLE


def cmd_serve(args: argparse.Namespace) -> None:
    data = Path(args.data).resolve()
    if args.init or not (data / "transport.json").is_file():
        data.mkdir(parents=True, exist_ok=True)
        transport = zkac.Keypair()
        issuer = zkac.BbsIssuer()
        pk = issuer.public_key()
        mg_rid = zkac.role_id(MGMT_ROLE)
        req = zkac.prepare_blind_request()
        sig = issuer.issue_blind(req.commitment_with_proof(), mg_rid, 1)
        _ = zkac.Credential.finalize(
            sig, req.member_secret(), req.prover_blind(), mg_rid, 1, pk
        )
        creds.save_server_transport(data / "transport.json", transport)
        creds.save_json(
            data / "mgmt_issuer.json",
            {"issuer_secret_b64": base64.b64encode(issuer.secret_key_bytes()).decode()},
        )
        m = creds.member_payload(sig, req, mg_rid, 1, pk)
        creds.save_json(data / "mgmt_member.json", m)
        print(f"Initialized server data in {data}", flush=True)
        print(
            "Copy mgmt_member.json for operators; distribute server_public_key from transport.json",
            flush=True,
        )
    relay = None if args.relay_port == 0 else args.relay_port
    relay_bind = args.relay_bind
    print(
        f"Starting: mgmt={args.mgmt_port} managed={args.managed_port} relay={relay!r} bind={relay_bind}",
        flush=True,
    )
    server_app.serve(data, args.mgmt_port, args.managed_port, relay, relay_bind)


def cmd_user_create(args: argparse.Namespace) -> None:
    uid = args.userid
    d = paths.ensure_user(uid)
    kp = zkac.Keypair()
    creds.save_transport_keypair(d / "transport.json", kp)
    creds.save_json(d / "profile.json", {"userid": uid})
    print(f"User {uid!r}: wrote {d / 'transport.json'}")


def cmd_registry_init(args: argparse.Namespace) -> None:
    base = paths.ensure_user(args.user)
    out = (base / "registries" / args.slug).resolve()
    r = registry_local.create_registry_bundle(args.slug, list(args.roles), out)
    print(json.dumps(r, indent=2))


def cmd_registry_push(args: argparse.Namespace) -> None:
    server_data = Path(args.server_data).resolve()
    reg_dir = (paths.user_dir(args.user) / "registries" / args.slug).resolve()
    reg = creds.load_json(reg_dir / "registry.json")
    mg = creds.load_json(server_data / "mgmt_member.json")
    t = creds.load_json(server_data / "transport.json")
    server_pk = zkac.PublicKey.from_bytes(base64.b64decode(t["server_public_key_b64"]))
    mgmt_cred = creds.load_member_credential(mg)
    cmd = {"cmd": "create_registry", "state_b64": reg["state_bytes_b64"], "state_cert_b64": reg["state_cert_b64"]}
    if args.update:
        cmd = {"cmd": "update_registry", "state_b64": reg["state_bytes_b64"], "state_cert_b64": reg["state_cert_b64"]}
    out = client_ops.mgmt_call(args.host, args.mgmt_port, server_pk, mgmt_cred, cmd)
    print(json.dumps(out, indent=2))


def cmd_connect(args: argparse.Namespace) -> None:
    reg_dir = (paths.user_dir(args.user) / "registries" / args.slug).resolve()
    reg_path = reg_dir / "registry.json"
    if reg_path.is_file():
        reg = creds.load_json(reg_path)
    elif args.registry_id:
        reg = {"registry_id_hex": args.registry_id}
    else:
        print(
            "Missing registry.json. Copy from the admin (public) or pass --registry-id.",
            file=sys.stderr,
        )
        raise SystemExit(1)
    cred_path = (reg_dir / args.credential).resolve()
    m = creds.load_json(cred_path)
    credential = creds.load_member_credential(m)
    server = creds.load_json(Path(args.server_file).expanduser().resolve())
    server_pk = zkac.PublicKey.from_bytes(base64.b64decode(server["server_public_key_b64"]))
    registry_id = bytes.fromhex(reg["registry_id_hex"])
    body = client_ops.managed_call(
        args.host,
        args.managed_port,
        server_pk,
        credential,
        registry_id,
        {"cmd": args.command},
    )
    print(json.dumps(body, indent=2))


def cmd_issuance_request(args: argparse.Namespace) -> None:
    """Queue a blind issuance request (payload is commitment bytes; localhost relay only)."""
    reg = creds.load_json(paths.user_dir(args.user) / "registries" / args.slug / "registry.json")
    issuance_pk = base64.b64decode(reg["issuance_public_key_b64"])
    req = zkac.prepare_blind_request()
    payload = req.commitment_with_proof()
    req_id = secrets.token_bytes(32)
    eph_pk = secrets.token_bytes(32)
    # Store locally for finalize
    pending = {
        "request_id_hex": req_id.hex(),
        "role_name": args.role,
        "member_secret_b64": base64.b64encode(req.member_secret()).decode(),
        "prover_blind_b64": base64.b64encode(req.prover_blind()).decode(),
        "issuer_public_key_b64": reg["admin_issuer_public_key_b64"],
        "epoch": 1,
    }
    pdir = paths.user_dir(args.user) / "pending"
    pdir.mkdir(parents=True, exist_ok=True)
    creds.save_json(pdir / f"{req_id.hex()}.json", pending)

    obj = {
        "cmd": "enqueue",
        "registry_id_hex": reg["registry_id_hex"],
        "role_name": args.role,
        "request_id_hex": req_id.hex(),
        "eph_pk_hex": eph_pk.hex(),
        "payload_b64": base64.b64encode(payload).decode(),
    }
    out = client_ops.relay_line(args.host, args.relay_port, obj)
    print(json.dumps(out, indent=2))
    print(f"Saved pending material to {pdir / (req_id.hex() + '.json')}")


def cmd_issuance_poll(args: argparse.Namespace) -> None:
    reg = creds.load_json(paths.user_dir(args.user) / "registries" / args.slug / "registry.json")
    obj = {
        "cmd": "poll",
        "registry_id_hex": reg["registry_id_hex"],
        "request_id_hex": args.request_id,
    }
    out = client_ops.relay_line(args.host, args.relay_port, obj)
    print(json.dumps(out, indent=2))
    if out.get("status") == "ready":
        pend = creds.load_json(paths.user_dir(args.user) / "pending" / f"{args.request_id}.json")
        pk = zkac.BbsPublicKey.from_bytes(base64.b64decode(pend["issuer_public_key_b64"]))
        rid = zkac.role_id(pend["role_name"])
        blind_sig = base64.b64decode(out["blind_sig_b64"])
        zkac.Credential.finalize(
            blind_sig,
            base64.b64decode(pend["member_secret_b64"]),
            base64.b64decode(pend["prover_blind_b64"]),
            rid,
            int(pend["epoch"]),
            pk,
        )
        out_dir = paths.user_dir(args.user) / "registries" / args.slug / "roles"
        out_dir.mkdir(parents=True, exist_ok=True)
        member = {
            "role_id_hex": rid.hex(),
            "epoch": int(pend["epoch"]),
            "blind_sig_b64": base64.b64encode(blind_sig).decode(),
            "member_secret_b64": pend["member_secret_b64"],
            "prover_blind_b64": pend["prover_blind_b64"],
            "issuer_public_key_b64": pend["issuer_public_key_b64"],
        }
        dest = out_dir / f"{pend['role_name']}.json"
        creds.save_json(dest, member)
        print(f"You have role {pend['role_name']!r}; credential saved to {dest}")
        print("Use: zkac-node connect ... --credential roles/{0}.json".format(pend["role_name"]))


def cmd_issuance_grant(args: argparse.Namespace) -> None:
    server_data = Path(args.server_data).resolve()
    reg_dir = (paths.user_dir(args.admin_user) / "registries" / args.slug).resolve()
    admin = creds.load_json(reg_dir / "admin.json")
    issuer = zkac.BbsIssuer.from_secret_key(base64.b64decode(admin["admin_issuer_secret_b64"]))
    t = creds.load_json(server_data / "transport.json")
    mg = creds.load_json(server_data / "mgmt_member.json")
    server_pk = zkac.PublicKey.from_bytes(base64.b64decode(t["server_public_key_b64"]))
    mgmt_cred = creds.load_member_credential(mg)

    peek = client_ops.mgmt_call(
        args.host,
        args.mgmt_port,
        server_pk,
        mgmt_cred,
        {"cmd": "issuance_peek", "registry_id_hex": args.registry_id},
    )
    pending = peek.get("pending") or []
    target = None
    for p in pending:
        if p["request_id_hex"] == args.request_id:
            target = p
            break
    if target is None:
        print("request_id not in pending queue", file=sys.stderr)
        raise SystemExit(1)
    commit = base64.b64decode(target["payload_b64"])
    role_id = bytes.fromhex(target["role_id_hex"])
    blind = issuer.issue_blind(commit, role_id, 1)
    out = client_ops.mgmt_call(
        args.host,
        args.mgmt_port,
        server_pk,
        mgmt_cred,
        {
            "cmd": "issuance_grant",
            "registry_id_hex": args.registry_id,
            "request_id_hex": args.request_id,
            "blind_sig_b64": base64.b64encode(blind).decode(),
        },
    )
    print(json.dumps(out, indent=2))


def cmd_issue_member_file(args: argparse.Namespace) -> None:
    """Admin issues a role credential locally (out-of-band handoff)."""
    reg_dir = (paths.user_dir(args.admin_user) / "registries" / args.slug).resolve()
    admin = creds.load_json(reg_dir / "admin.json")
    issuer = zkac.BbsIssuer.from_secret_key(base64.b64decode(admin["admin_issuer_secret_b64"]))
    pk = issuer.public_key()
    rid = zkac.role_id(args.role)
    req = zkac.prepare_blind_request()
    sig = issuer.issue_blind(req.commitment_with_proof(), rid, 1)
    payload = creds.member_payload(sig, req, rid, 1, pk)
    out_path = reg_dir / "issued" / f"{args.role}_{args.target_user}.json"
    creds.save_json(out_path, payload)
    print(f"Wrote {out_path}")


def cmd_registry_import_public(args: argparse.Namespace) -> None:
    """Copy public registry.json (state + cert summary) from another local user."""
    src = (paths.user_dir(args.from_user) / "registries" / args.slug / "registry.json").resolve()
    dst_dir = paths.ensure_user(args.to_user) / "registries" / args.slug
    dst_dir.mkdir(parents=True, exist_ok=True)
    shutil.copy(src, dst_dir / "registry.json")
    print(f"Copied {src} -> {dst_dir / 'registry.json'}")


def cmd_import_credential(args: argparse.Namespace) -> None:
    src = Path(args.file).resolve()
    data = creds.load_json(src)
    dest = paths.user_dir(args.user) / "registries" / args.slug / "roles" / f"{args.name}.json"
    creds.save_json(dest, data)
    print(f"Imported credential to {dest}")


def cmd_pin_server(args: argparse.Namespace) -> None:
    """Save server public key + ports for a user (from server's transport.json)."""
    tpath = Path(args.transport_file).resolve()
    t = creds.load_json(tpath)
    d = paths.ensure_user(args.user) / "servers" / args.name
    d.mkdir(parents=True, exist_ok=True)
    creds.save_json(
        d / "server.json",
        {
            "server_public_key_b64": t["server_public_key_b64"],
            "host": args.host,
            "mgmt_port": args.mgmt_port,
            "managed_port": args.managed_port,
            "relay_port": args.relay_port,
        },
    )
    print(f"Pinned server {args.name!r} for user {args.user!r} -> {d / 'server.json'}")


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="zkac-node",
        description="ZKAC node CLI (server + client). "
        "Requires the 'zkac' package from this repo (maturin develop). "
        "Set ZKAC_HOME to override ~/.zkac/. "
        "Issuance relay queues are in-memory — use the same server process between request and grant.",
    )
    sub = p.add_subparsers(dest="command", required=True)

    ps = sub.add_parser("serve", help="Run server (registry-capable node)")
    ps.add_argument("--data", type=Path, required=True, help="Server data directory")
    ps.add_argument("--init", action="store_true", help="Initialize data dir (transport + mgmt issuer)")
    ps.add_argument("--mgmt-port", type=int, default=7400)
    ps.add_argument("--managed-port", type=int, default=7401)
    ps.add_argument(
        "--relay-port",
        type=int,
        default=7402,
        help="Plaintext relay for issuance queue (localhost only). Use 0 to disable.",
    )
    ps.add_argument("--relay-bind", default="127.0.0.1")
    ps.set_defaults(func=cmd_serve)

    su = sub.add_parser("user-create", help="Create a user directory under ~/.zkac/")
    su.add_argument("userid")
    su.set_defaults(func=cmd_user_create)

    pi = sub.add_parser("pin-server", help="Save server connection info for a user")
    pi.add_argument("--user", required=True)
    pi.add_argument("--name", required=True, help="Label for this server")
    pi.add_argument("--transport-file", required=True, help="Path to server's transport.json")
    pi.add_argument("--host", required=True)
    pi.add_argument("--mgmt-port", type=int, default=7400)
    pi.add_argument("--managed-port", type=int, default=7401)
    pi.add_argument("--relay-port", type=int, default=7402)
    pi.set_defaults(func=cmd_pin_server)

    ri = sub.add_parser("registry-init", help="Create a client-managed registry offline")
    ri.add_argument("--user", required=True)
    ri.add_argument("--slug", required=True)
    ri.add_argument("--roles", nargs="+", required=True)
    ri.set_defaults(func=cmd_registry_init)

    rp = sub.add_parser("registry-push", help="Push registry snapshot to server (mgmt channel)")
    rp.add_argument("--user", required=True)
    rp.add_argument("--slug", required=True)
    rp.add_argument("--server-data", type=Path, required=True, help="Server data dir (transport + mgmt_member)")
    rp.add_argument("--host", default="127.0.0.1")
    rp.add_argument("--mgmt-port", type=int, default=7400)
    rp.add_argument("--update", action="store_true", help="Send update_registry instead of create")
    rp.set_defaults(func=cmd_registry_push)

    cc = sub.add_parser("connect", help="Managed ZKAC session test (whoami/get_registry)")
    cc.add_argument("--user", required=True)
    cc.add_argument("--slug", required=True)
    cc.add_argument("--credential", default="roles/member.json", help="Relative path under registries/<slug>/")
    cc.add_argument(
        "--registry-id",
        default="",
        help="If registry.json is absent, hex registry id (pin public registry first)",
    )
    cc.add_argument("--server-file", required=True, help="Pinned server.json from pin-server")
    cc.add_argument("--host", required=True)
    cc.add_argument("--managed-port", type=int, default=7401)
    cc.add_argument("--command", default="whoami")
    cc.set_defaults(func=cmd_connect)

    ir = sub.add_parser(
        "issuance-request",
        help="Enqueue blind commitment on relay (localhost; see help text)",
    )
    ir.add_argument("--user", required=True)
    ir.add_argument("--slug", required=True)
    ir.add_argument("--role", required=True)
    ir.add_argument("--host", default="127.0.0.1")
    ir.add_argument("--relay-port", type=int, default=7402)
    ir.set_defaults(func=cmd_issuance_request)

    ip = sub.add_parser("issuance-poll", help="Poll relay for blind signature")
    ip.add_argument("--user", required=True)
    ip.add_argument("--slug", required=True)
    ip.add_argument("--request-id", required=True, dest="request_id")
    ip.add_argument("--host", default="127.0.0.1")
    ip.add_argument("--relay-port", type=int, default=7402)
    ip.set_defaults(func=cmd_issuance_poll)

    ig = sub.add_parser("issuance-grant", help="Admin: issue blind sig and push grant to server")
    ig.add_argument("--admin-user", required=True)
    ig.add_argument("--slug", required=True)
    ig.add_argument("--server-data", type=Path, required=True)
    ig.add_argument("--registry-id", required=True, dest="registry_id")
    ig.add_argument("--request-id", required=True, dest="request_id")
    ig.add_argument("--host", default="127.0.0.1")
    ig.add_argument("--mgmt-port", type=int, default=7400)
    ig.set_defaults(func=cmd_issuance_grant)

    im = sub.add_parser("issue-member", help="Admin: issue credential to file for handoff")
    im.add_argument("--admin-user", required=True)
    im.add_argument("--slug", required=True)
    im.add_argument("--role", required=True)
    im.add_argument("--target-user", required=True)
    im.set_defaults(func=cmd_issue_member_file)

    irp = sub.add_parser(
        "registry-import-public",
        help="Copy registry.json from another user (public state + cert metadata)",
    )
    irp.add_argument("--from-user", required=True, dest="from_user")
    irp.add_argument("--to-user", required=True, dest="to_user")
    irp.add_argument("--slug", required=True)
    irp.set_defaults(func=cmd_registry_import_public)

    ii = sub.add_parser("import-credential", help="Copy a member json into a user's registry folder")
    ii.add_argument("--user", required=True)
    ii.add_argument("--slug", required=True)
    ii.add_argument("--name", required=True, help="Filename base (e.g. analyst)")
    ii.add_argument("file", type=Path)
    ii.set_defaults(func=cmd_import_credential)

    return p


def main() -> None:
    args = build_parser().parse_args()
    args.func(args)


if __name__ == "__main__":
    main()