ZKAC/demo/file_share_server.py
everbarry d5ae07973a polish and self-contain file-share demo UI
Co-authored-by: Cursor <cursoragent@cursor.com>
2026-05-07 18:39:39 +02:00

654 lines
24 KiB
Python

#!/usr/bin/env python3
"""
ZKAC opaque file-share server (demo).
Headless TCP service that combines:
* The same management-channel wire protocol as ``zkac-node serve`` so existing
``zkac-node registry create/get/update`` commands work unchanged.
* A new file-share channel that, after BBS+ role authentication, exposes
bucket primitives. The server only ever sees opaque ciphertext blobs and
per-recipient encrypted role grants; file names, contents and per-role
visibility masks are never visible server-side.
Run with::
uv run python demo/file_share_server.py --host 127.0.0.1 --port 9879
The server transport public key is printed at start-up; pin it on each client
with ``zkac-node server pin <userid> 127.0.0.1:9879 --key <hex>``.
"""
from __future__ import annotations
import argparse
import base64
import json
import os
import socket
import sys
import threading
import traceback
import uuid
from pathlib import Path
from typing import Any
import zkac
from zkac.tcp import FramedSession, server_handshake_anon
# ── small helpers ────────────────────────────────────────────────────
def _b64(data: bytes) -> str:
return base64.b64encode(data).decode("ascii")
def _unb64(s: str) -> bytes:
return base64.b64decode(s)
def _chmod(path: Path, mode: int) -> None:
try:
os.chmod(path, mode)
except OSError:
pass
def _write_private_json(path: Path, payload: dict) -> None:
path.write_text(json.dumps(payload, indent=2))
_chmod(path, 0o600)
def _is_loopback(host: str) -> bool:
return host.strip().lower() in {"127.0.0.1", "::1", "localhost"}
def _safe_id(value: str) -> str:
"""Strict allowlist on identifiers used as filesystem names."""
if not isinstance(value, str) or not value:
raise ValueError("missing identifier")
if len(value) > 128 or any(c not in "0123456789abcdefABCDEF-" for c in value):
raise ValueError(f"invalid identifier {value!r}")
return value
# ── opaque on-disk store ─────────────────────────────────────────────
class _FileShareStore:
"""Persists registry snapshots and opaque bucket state under one data dir."""
def __init__(self, data_dir: Path) -> None:
self._dir = data_dir
self._reg_dir = data_dir / "registries"
self._buckets_dir = data_dir / "buckets"
for d in (self._dir, self._reg_dir, self._buckets_dir):
d.mkdir(parents=True, exist_ok=True)
_chmod(d, 0o700)
self._lock = threading.Lock()
# transport key
def load_or_create_keypair(self) -> zkac.Keypair:
kf = self._dir / "server_key.json"
if kf.exists():
data = json.loads(kf.read_text())
return zkac.Keypair.from_secret_key(_unb64(data["secret_b64"]))
kp = zkac.Keypair()
_write_private_json(kf, {
"secret_b64": _b64(kp.secret_key_bytes()),
"public_b64": _b64(kp.public_key().to_bytes()),
})
return kp
# registries (mirrors zkac-node serve so the CLI works unchanged)
def save_registry(self, rid_hex: str, state_bytes: bytes, cert_bytes: bytes) -> None:
with self._lock:
(self._reg_dir / f"{rid_hex}.state").write_bytes(state_bytes)
(self._reg_dir / f"{rid_hex}.cert").write_bytes(cert_bytes)
def load_all_registries(self, mgr: zkac.RegistryManager) -> int:
n = 0
for p in sorted(self._reg_dir.glob("*.state")):
cert = self._reg_dir / f"{p.stem}.cert"
if not cert.exists():
continue
try:
mgr.restore(p.read_bytes(), cert.read_bytes())
n += 1
except Exception as exc:
print(f"[fs-server] skip registry {p.stem}: {exc}")
return n
# buckets
def _bucket_dir(self, bucket_id: str) -> Path:
return self._buckets_dir / _safe_id(bucket_id)
def _bucket_meta_path(self, bucket_id: str) -> Path:
return self._bucket_dir(bucket_id) / "meta.json"
def bucket_create(self, bucket_id: str, owner_registry_id_hex: str) -> None:
bd = self._bucket_dir(bucket_id)
with self._lock:
if bd.exists():
raise RuntimeError("bucket already exists")
(bd / "blobs").mkdir(parents=True)
(bd / "role_grants").mkdir()
_chmod(bd, 0o700)
_write_private_json(self._bucket_meta_path(bucket_id), {
"bucket_id": bucket_id,
"owner_registry_id": owner_registry_id_hex,
"finalized": False,
"role_acl": {}, # role_id_hex -> {"version": int, "allowed_blob_ids": [blob_id...]}
})
def bucket_meta(self, bucket_id: str) -> dict:
return json.loads(self._bucket_meta_path(bucket_id).read_text())
def _bucket_require_owner(self, bucket_id: str, registry_id_hex: str) -> dict:
meta = self.bucket_meta(bucket_id)
if meta.get("owner_registry_id") != registry_id_hex:
raise RuntimeError("not bucket owner")
return meta
def bucket_set_finalized(self, bucket_id: str, registry_id_hex: str, finalized: bool) -> None:
with self._lock:
meta = self._bucket_require_owner(bucket_id, registry_id_hex)
meta["finalized"] = bool(finalized)
_write_private_json(self._bucket_meta_path(bucket_id), meta)
def bucket_delete(self, bucket_id: str, registry_id_hex: str) -> None:
with self._lock:
self._bucket_require_owner(bucket_id, registry_id_hex)
bd = self._bucket_dir(bucket_id)
for sub in ("blobs", "role_grants"):
d = bd / sub
if d.is_dir():
for p in d.iterdir():
try:
p.unlink()
except OSError:
pass
try:
d.rmdir()
except OSError:
pass
try:
self._bucket_meta_path(bucket_id).unlink()
except OSError:
pass
try:
bd.rmdir()
except OSError:
pass
def bucket_put_blob(
self,
bucket_id: str,
registry_id_hex: str,
blob_id: str,
ciphertext: bytes,
) -> None:
with self._lock:
self._bucket_require_owner(bucket_id, registry_id_hex)
(self._bucket_dir(bucket_id) / "blobs" / _safe_id(blob_id)).write_bytes(ciphertext)
def bucket_put_role_grant(
self,
bucket_id: str,
registry_id_hex: str,
recipient_pk_hex: str,
role_id_hex: str,
acl_version: int,
eph_pk_b64: str,
ciphertext_b64: str,
) -> None:
with self._lock:
self._bucket_require_owner(bucket_id, registry_id_hex)
_safe_id(recipient_pk_hex)
payload = {
"bucket_id": bucket_id,
"role_id_hex": _safe_id(role_id_hex),
"acl_version": int(acl_version),
"eph_pk_b64": eph_pk_b64,
"ciphertext_b64": ciphertext_b64,
}
target = self._bucket_dir(bucket_id) / "role_grants" / f"{recipient_pk_hex}.json"
_write_private_json(target, payload)
def bucket_set_role_acl(
self,
bucket_id: str,
registry_id_hex: str,
role_id_hex: str,
allowed_blob_ids: list[str],
) -> None:
with self._lock:
meta = self._bucket_require_owner(bucket_id, registry_id_hex)
role_acl = dict(meta.get("role_acl", {}))
key = _safe_id(role_id_hex)
prev = role_acl.get(key, {})
prev_version = int(prev.get("version", 0)) if isinstance(prev, dict) else 0
role_acl[key] = {
"version": prev_version + 1,
"allowed_blob_ids": [_safe_id(b) for b in allowed_blob_ids],
}
meta["role_acl"] = role_acl
_write_private_json(self._bucket_meta_path(bucket_id), meta)
def bucket_get_blob(self, bucket_id: str, blob_id: str) -> bytes:
return (self._bucket_dir(bucket_id) / "blobs" / _safe_id(blob_id)).read_bytes()
def bucket_blob_allowed_for_role(self, bucket_id: str, role_id_hex: str, blob_id: str) -> bool:
meta = self.bucket_meta(bucket_id)
role_acl = meta.get("role_acl", {})
role_meta = role_acl.get(_safe_id(role_id_hex))
if not isinstance(role_meta, dict):
return False
allowed = role_meta.get("allowed_blob_ids")
if not isinstance(allowed, list):
return False
return _safe_id(blob_id) in {str(v) for v in allowed}
def bucket_role_acl(self, bucket_id: str, role_id_hex: str) -> dict:
meta = self.bucket_meta(bucket_id)
role_acl = meta.get("role_acl", {})
role_meta = role_acl.get(_safe_id(role_id_hex), {})
if not isinstance(role_meta, dict):
return {"version": 0, "allowed_blob_ids": []}
version = int(role_meta.get("version", 0))
allowed = role_meta.get("allowed_blob_ids", [])
if not isinstance(allowed, list):
allowed = []
return {
"version": version,
"allowed_blob_ids": [_safe_id(str(v)) for v in allowed],
}
def bucket_get_role_grant(self, bucket_id: str, recipient_pk_hex: str) -> dict:
path = self._bucket_dir(bucket_id) / "role_grants" / f"{_safe_id(recipient_pk_hex)}.json"
return json.loads(path.read_text())
def buckets_for_recipient(self, recipient_pk_hex: str) -> list[str]:
"""Bucket ids that have an encrypted role grant addressed to ``recipient_pk_hex``."""
_safe_id(recipient_pk_hex)
out: list[str] = []
for bd in sorted(self._buckets_dir.iterdir()):
grant = bd / "role_grants" / f"{recipient_pk_hex}.json"
if grant.is_file():
out.append(bd.name)
return out
def buckets_owned_by(self, registry_id_hex: str) -> list[str]:
out: list[str] = []
for bd in sorted(self._buckets_dir.iterdir()):
meta = bd / "meta.json"
if not meta.is_file():
continue
try:
if json.loads(meta.read_text()).get("owner_registry_id") == registry_id_hex:
out.append(bd.name)
except (OSError, json.JSONDecodeError):
continue
return out
# ── command dispatch (inside encrypted session) ──────────────────────
def _dispatch_mgmt(
cmd: dict,
mgr: zkac.RegistryManager,
store: _FileShareStore,
server_pk_b64: str,
transcript_hash: bytes,
) -> dict:
"""Registry management commands, wire-compatible with ``zkac-node`` CLI."""
try:
action = cmd.get("cmd")
rid_hex = cmd.get("auth_registry_id")
admin_proof_b64 = cmd.get("admin_proof_b64")
def _require_admin_for(target_rid_hex: str) -> None:
if rid_hex != target_rid_hex:
raise RuntimeError("auth_registry_id must match command registry_id")
if not isinstance(admin_proof_b64, str) or not admin_proof_b64:
raise RuntimeError("missing admin_proof_b64")
if not mgr.verify_admin(
bytes.fromhex(target_rid_hex),
_unb64(admin_proof_b64),
transcript_hash,
):
raise RuntimeError("admin authorization failed")
if action == "server_info":
return {"ok": True, "server_public_key_b64": server_pk_b64}
if action == "create_registry":
state_bytes = _unb64(cmd["state_bytes_b64"])
state_cert = _unb64(cmd["state_cert_b64"])
auth_rid = cmd.get("auth_registry_id")
if not isinstance(auth_rid, str):
raise RuntimeError("missing auth_registry_id")
if not isinstance(admin_proof_b64, str) or not admin_proof_b64:
raise RuntimeError("missing admin_proof_b64")
tmp_mgr = zkac.RegistryManager()
expected_rid = tmp_mgr.create(state_bytes, state_cert).hex()
if expected_rid != auth_rid:
raise RuntimeError("auth_registry_id does not match certified state")
if not tmp_mgr.verify_admin(
bytes.fromhex(expected_rid),
_unb64(admin_proof_b64),
transcript_hash,
):
raise RuntimeError("admin authorization failed for create_registry")
rid = mgr.create(state_bytes, state_cert)
store.save_registry(rid.hex(), state_bytes, state_cert)
return {"ok": True, "registry_id": rid.hex()}
if action == "get_registry":
rid_hex_cmd = cmd["registry_id"]
_require_admin_for(rid_hex_cmd)
rid = bytes.fromhex(rid_hex_cmd)
state_bytes, state_cert = mgr.get(rid)
return {
"ok": True,
"state_bytes_b64": _b64(state_bytes),
"state_cert_b64": _b64(state_cert),
}
if action == "update_registry":
rid_hex_cmd = cmd["registry_id"]
_require_admin_for(rid_hex_cmd)
rid = bytes.fromhex(rid_hex_cmd)
state_bytes = _unb64(cmd["state_bytes_b64"])
state_cert = _unb64(cmd["state_cert_b64"])
mgr.update(rid, state_bytes, state_cert)
store.save_registry(rid_hex_cmd, state_bytes, state_cert)
return {"ok": True}
return {"error": f"unknown command: {action}"}
except Exception as exc:
return {"error": str(exc)}
def _dispatch_fs(
cmd: dict,
store: _FileShareStore,
ctx: dict,
) -> dict:
"""File-share commands authenticated via ``ctx`` (registry_id, role_id, issuance_pk_hex)."""
try:
action = cmd.get("cmd")
registry_id_hex: str = ctx["registry_id_hex"]
role_id: bytes = ctx["role_id"]
issuance_pk_hex: str = ctx["issuance_pk_hex"]
is_admin = role_id == zkac.admin_role_id()
def _require_admin() -> None:
if not is_admin:
raise RuntimeError("admin role required for this command")
if action == "whoami":
return {
"ok": True,
"registry_id": registry_id_hex,
"role_id": role_id.hex(),
"is_admin": is_admin,
"issuance_pk_hex": issuance_pk_hex,
}
if action == "bucket_create":
_require_admin()
bid = cmd.get("bucket_id") or uuid.uuid4().hex
store.bucket_create(bid, registry_id_hex)
return {"ok": True, "bucket_id": bid}
if action == "bucket_put_blob":
_require_admin()
bid = cmd["bucket_id"]
blob_id = cmd["blob_id"]
ciphertext = _unb64(cmd["ciphertext_b64"])
store.bucket_put_blob(bid, registry_id_hex, blob_id, ciphertext)
return {"ok": True}
if action == "bucket_put_role_grant":
_require_admin()
store.bucket_put_role_grant(
cmd["bucket_id"],
registry_id_hex,
cmd["recipient_pk_hex"],
cmd["role_id_hex"],
int(cmd["acl_version"]),
cmd["eph_pk_b64"],
cmd["ciphertext_b64"],
)
return {"ok": True}
if action == "bucket_set_role_acl":
_require_admin()
store.bucket_set_role_acl(
cmd["bucket_id"],
registry_id_hex,
cmd["role_id_hex"],
list(cmd.get("allowed_blob_ids", [])),
)
return {"ok": True}
if action == "bucket_get_role_acl":
_require_admin()
role_acl = store.bucket_role_acl(cmd["bucket_id"], cmd["role_id_hex"])
return {"ok": True, **role_acl}
if action == "bucket_finalize":
_require_admin()
store.bucket_set_finalized(cmd["bucket_id"], registry_id_hex, True)
return {"ok": True}
if action == "bucket_delete":
_require_admin()
store.bucket_delete(cmd["bucket_id"], registry_id_hex)
return {"ok": True}
if action == "bucket_list_owned":
_require_admin()
return {"ok": True, "bucket_ids": store.buckets_owned_by(registry_id_hex)}
if action == "fs_buckets":
return {"ok": True, "bucket_ids": store.buckets_for_recipient(issuance_pk_hex)}
if action == "fs_get_role_grant":
bid = cmd["bucket_id"]
grant = store.bucket_get_role_grant(bid, issuance_pk_hex)
if not is_admin:
expected_role = _safe_id(role_id.hex())
granted_role_raw = grant.get("role_id_hex")
if not isinstance(granted_role_raw, str) or not granted_role_raw:
raise RuntimeError("permissions updated; role grant is outdated, request a fresh grant")
granted_role = _safe_id(granted_role_raw)
if granted_role != expected_role:
raise RuntimeError("role grant not valid for authenticated role")
current_acl = store.bucket_role_acl(bid, expected_role)
acl_version = grant.get("acl_version")
if not isinstance(acl_version, int):
raise RuntimeError("permissions updated; role grant is outdated, request a fresh grant")
if acl_version != int(current_acl.get("version", -2)):
raise RuntimeError("permissions updated; request a fresh role grant")
return {"ok": True, **grant}
if action == "fs_get_blob":
bid = cmd["bucket_id"]
blob_id = cmd["blob_id"]
if not is_admin and not store.bucket_blob_allowed_for_role(bid, role_id.hex(), blob_id):
raise RuntimeError("blob access denied by role mask")
data = store.bucket_get_blob(bid, blob_id)
return {"ok": True, "ciphertext_b64": _b64(data)}
return {"error": f"unknown command: {action}"}
except Exception as exc:
return {"error": str(exc)}
# ── per-connection handler ────────────────────────────────────────────
def _handle_conn(
conn: socket.socket,
addr: tuple,
node: zkac.Node,
mgr: zkac.RegistryManager,
store: _FileShareStore,
server_pk_b64: str,
idle_timeout_s: float,
slots: threading.BoundedSemaphore,
) -> None:
peer = f"{addr[0]}:{addr[1]}"
try:
conn.settimeout(idle_timeout_s)
session = server_handshake_anon(conn, node)
framed = FramedSession(conn, session)
transcript_hash = bytes(session.transcript_hash())
hello = json.loads(framed.recv())
op = hello.get("op")
if op == "mgmt":
while True:
try:
cmd = json.loads(framed.recv())
except (ConnectionError, OSError):
break
resp = _dispatch_mgmt(cmd, mgr, store, server_pk_b64, transcript_hash)
framed.send(json.dumps(resp).encode())
return
if op == "fs":
try:
registry_id = bytes.fromhex(hello["registry_id"])
role_id = bytes.fromhex(hello["role_id"])
proof = _unb64(hello["bbs_auth_b64"])
issuance_pk_hex = hello["issuance_pk_hex"]
_safe_id(issuance_pk_hex)
except (KeyError, ValueError) as exc:
framed.send(json.dumps({"error": f"invalid fs hello: {exc}"}).encode())
return
if role_id == zkac.admin_role_id():
ok = mgr.verify_admin(registry_id, proof, transcript_hash)
else:
ok = mgr.verify_presentation(registry_id, role_id, proof, transcript_hash)
if not ok:
framed.send(json.dumps({"error": "auth failed"}).encode())
return
framed.send(json.dumps({"ok": True, "status": "authenticated"}).encode())
ctx = {
"registry_id_hex": registry_id.hex(),
"role_id": role_id,
"issuance_pk_hex": issuance_pk_hex,
}
while True:
try:
cmd = json.loads(framed.recv())
except (ConnectionError, OSError):
break
resp = _dispatch_fs(cmd, store, ctx)
framed.send(json.dumps(resp).encode())
return
framed.send(json.dumps({"error": f"unknown op: {op}"}).encode())
except (ConnectionError, BrokenPipeError, OSError):
pass
except Exception as exc:
print(f"[fs-server] {peer} error: {exc}")
traceback.print_exc()
finally:
conn.close()
slots.release()
# ── entry point ──────────────────────────────────────────────────────
def serve(
data_dir: Path,
host: str = "127.0.0.1",
port: int = 9879,
*,
max_connections: int = 64,
idle_timeout_s: float = 60.0,
listen_backlog: int = 64,
allow_non_loopback: bool = False,
) -> None:
data_dir.mkdir(parents=True, exist_ok=True)
store = _FileShareStore(data_dir)
kp = store.load_or_create_keypair()
server_pk_b64 = _b64(kp.public_key().to_bytes())
pk_hex = kp.public_key().to_bytes().hex()
node = zkac.Node(kp)
mgr = zkac.RegistryManager()
n = store.load_all_registries(mgr)
print(f"[fs-server] data dir: {data_dir}")
print(f"[fs-server] server transport public key (pin OOB): {pk_hex}")
print(f"[fs-server] loaded {n} registries")
print(f"[fs-server] listening on {host}:{port}")
if not _is_loopback(host) and not allow_non_loopback:
raise RuntimeError(
"refusing to bind outside loopback. "
"Use --allow-non-loopback only when exposure is intentional."
)
if not _is_loopback(host):
print(f"[fs-server] warning: binding outside loopback: {host}:{port}", file=sys.stderr)
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock.bind((host, port))
slots = threading.BoundedSemaphore(max_connections)
sock.listen(listen_backlog)
try:
while True:
conn, addr = sock.accept()
if not slots.acquire(blocking=False):
conn.close()
continue
threading.Thread(
target=_handle_conn,
args=(conn, addr, node, mgr, store, server_pk_b64, idle_timeout_s, slots),
daemon=True,
).start()
except KeyboardInterrupt:
print("\n[fs-server] shutdown")
finally:
sock.close()
def main() -> None:
p = argparse.ArgumentParser(description="ZKAC opaque file-share server (demo)")
p.add_argument("--host", default="127.0.0.1")
p.add_argument("--port", type=int, default=9879)
p.add_argument(
"--data-dir",
type=Path,
default=Path(__file__).resolve().parent / "fs_data",
help="server state (transport key, registries, opaque buckets)",
)
p.add_argument("--idle-timeout", type=float, default=60.0)
p.add_argument("--max-connections", type=int, default=64)
p.add_argument("--allow-non-loopback", action="store_true")
args = p.parse_args()
serve(
args.data_dir,
host=args.host,
port=args.port,
max_connections=args.max_connections,
idle_timeout_s=args.idle_timeout,
allow_non_loopback=args.allow_non_loopback,
)
if __name__ == "__main__":
main()