quicproquo/examples/python/qpc_client.py

#!/usr/bin/env python3
"""quicprochat Python client -- ctypes wrapper around libquicprochat_ffi.

Usage:
    python qpq_client.py --server 127.0.0.1:7000 \\
        --ca-cert server-cert.der \\
        --server-name localhost \\
        --username alice --password secret \\
        [--send-to bob --message "hello"] \\
        [--receive --timeout 5000]
"""

from __future__ import annotations

import argparse
import ctypes
import json
import os
import sys
from ctypes import (
    POINTER,
    c_char_p,
    c_int,
    c_size_t,
    c_uint32,
    c_void_p,
)
from pathlib import Path

# ---------------------------------------------------------------------------
# Status codes (must match quicprochat-ffi/src/lib.rs)
# ---------------------------------------------------------------------------
QPQ_OK = 0
QPQ_ERROR = 1
QPQ_AUTH_FAILED = 2
QPQ_TIMEOUT = 3
QPQ_NOT_CONNECTED = 4

STATUS_NAMES = {
    QPQ_OK: "OK",
    QPQ_ERROR: "ERROR",
    QPQ_AUTH_FAILED: "AUTH_FAILED",
    QPQ_TIMEOUT: "TIMEOUT",
    QPQ_NOT_CONNECTED: "NOT_CONNECTED",
}


def _status_name(code: int) -> str:
    return STATUS_NAMES.get(code, f"UNKNOWN({code})")


# ---------------------------------------------------------------------------
# Library loading
# ---------------------------------------------------------------------------

def _find_library() -> str:
    """Locate libquicprochat_ffi shared library."""
    env = os.environ.get("QPQ_FFI_LIB")
    if env:
        return env

    # Walk up from this script to find the cargo target directory.
    repo_root = Path(__file__).resolve().parent.parent.parent
    for profile in ("release", "debug"):
        for ext in ("so", "dylib", "dll"):
            candidate = repo_root / "target" / profile / f"libquicprochat_ffi.{ext}"
            if candidate.exists():
                return str(candidate)

    print(
        "ERROR: Could not find libquicprochat_ffi. "
        "Build with: cargo build --release -p quicprochat-ffi\n"
        "Or set QPQ_FFI_LIB=/path/to/libquicprochat_ffi.so",
        file=sys.stderr,
    )
    sys.exit(1)


def _load_library(path: str) -> ctypes.CDLL:
    """Load the FFI library and declare function signatures."""
    lib = ctypes.cdll.LoadLibrary(path)

    # qpq_connect(server, ca_cert, server_name) -> *mut QpqHandle
    lib.qpq_connect.argtypes = [c_char_p, c_char_p, c_char_p]
    lib.qpq_connect.restype = c_void_p

    # qpq_login(handle, username, password) -> i32
    lib.qpq_login.argtypes = [c_void_p, c_char_p, c_char_p]
    lib.qpq_login.restype = c_int

    # qpq_send(handle, recipient, message, message_len) -> i32
    lib.qpq_send.argtypes = [c_void_p, c_char_p, ctypes.POINTER(ctypes.c_uint8), c_size_t]
    lib.qpq_send.restype = c_int

    # qpq_receive(handle, timeout_ms, out_json) -> i32
    lib.qpq_receive.argtypes = [c_void_p, c_uint32, POINTER(c_char_p)]
    lib.qpq_receive.restype = c_int

    # qpq_disconnect(handle)
    lib.qpq_disconnect.argtypes = [c_void_p]
    lib.qpq_disconnect.restype = None

    # qpq_last_error(handle) -> *const c_char
    lib.qpq_last_error.argtypes = [c_void_p]
    lib.qpq_last_error.restype = c_char_p

    # qpq_free_string(ptr)
    lib.qpq_free_string.argtypes = [c_char_p]
    lib.qpq_free_string.restype = None

    return lib


# ---------------------------------------------------------------------------
# High-level wrapper
# ---------------------------------------------------------------------------

class QpqClient:
    """Synchronous Python client wrapping the quicprochat C FFI."""

    def __init__(self, lib: ctypes.CDLL):
        self._lib = lib
        self._handle: int | None = None

    def connect(self, server: str, ca_cert: str, server_name: str) -> None:
        handle = self._lib.qpq_connect(
            server.encode(), ca_cert.encode(), server_name.encode(),
        )
        if not handle:
            raise ConnectionError("qpq_connect returned null -- check server address and CA cert")
        self._handle = handle

    def login(self, username: str, password: str) -> None:
        rc = self._lib.qpq_login(self._handle, username.encode(), password.encode())
        if rc != QPQ_OK:
            err = self._last_error()
            raise RuntimeError(f"qpq_login failed ({_status_name(rc)}): {err}")

    def send(self, recipient: str, message: str) -> None:
        msg_bytes = message.encode()
        buf = (ctypes.c_uint8 * len(msg_bytes))(*msg_bytes)
        rc = self._lib.qpq_send(self._handle, recipient.encode(), buf, len(msg_bytes))
        if rc != QPQ_OK:
            err = self._last_error()
            raise RuntimeError(f"qpq_send failed ({_status_name(rc)}): {err}")

    def receive(self, timeout_ms: int = 5000) -> list[str]:
        out = c_char_p()
        rc = self._lib.qpq_receive(self._handle, c_uint32(timeout_ms), ctypes.byref(out))
        if rc == QPQ_TIMEOUT:
            return []
        if rc != QPQ_OK:
            err = self._last_error()
            raise RuntimeError(f"qpq_receive failed ({_status_name(rc)}): {err}")

        result: list[str] = []
        if out.value:
            result = json.loads(out.value.decode())
            self._lib.qpq_free_string(out)
        return result

    def disconnect(self) -> None:
        if self._handle:
            self._lib.qpq_disconnect(self._handle)
            self._handle = None

    def _last_error(self) -> str:
        err = self._lib.qpq_last_error(self._handle)
        if err:
            return err.decode(errors="replace")
        return "(no error message)"

    def __enter__(self):
        return self

    def __exit__(self, *_):
        self.disconnect()


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(description="quicprochat Python FFI client")
    parser.add_argument("--server", required=True, help="Server address (host:port)")
    parser.add_argument("--ca-cert", required=True, help="Path to CA certificate (DER)")
    parser.add_argument("--server-name", default="localhost", help="TLS server name")
    parser.add_argument("--username", required=True, help="OPAQUE username")
    parser.add_argument("--password", required=True, help="OPAQUE password")
    parser.add_argument("--send-to", help="Recipient username for sending a message")
    parser.add_argument("--message", help="Message text to send")
    parser.add_argument("--receive", action="store_true", help="Receive pending messages")
    parser.add_argument("--timeout", type=int, default=5000, help="Receive timeout (ms)")
    args = parser.parse_args()

    lib_path = _find_library()
    print(f"Using library: {lib_path}")
    lib = _load_library(lib_path)

    with QpqClient(lib) as client:
        print(f"Connecting to {args.server}...")
        client.connect(args.server, args.ca_cert, args.server_name)
        print("Connected.")

        print(f"Logging in as {args.username}...")
        client.login(args.username, args.password)
        print("Logged in.")

        if args.send_to and args.message:
            print(f"Sending to {args.send_to}: {args.message}")
            client.send(args.send_to, args.message)
            print("Sent.")

        if args.receive:
            print(f"Receiving (timeout={args.timeout}ms)...")
            messages = client.receive(timeout_ms=args.timeout)
            if messages:
                print(f"Received {len(messages)} message(s):")
                for i, msg in enumerate(messages, 1):
                    print(f"  [{i}] {msg}")
            else:
                print("No messages received.")

        print("Disconnecting...")

    print("Done.")


if __name__ == "__main__":
    main()