kx-auth-cli 0.5.0b1__py3-none-any.whl

kx_auth_cli/__init__.py

@@ -0,0 +1,7 @@
+"""``kx auth`` — the agent-facing auth CLI for the kx-mcp control plane.
+
+A thin ``kx`` umbrella (stdlib argparse) whose ``auth`` group exposes the auth flows as a headless
+surface for an autonomous MCP client: structured ``--json`` output and a stable exit-code contract so
+the agent branches on codes, not prose. Subcommands: ``introspect``, ``login``, ``exchange``,
+``assert``. Verification reuses the lean :mod:`kx_auth_core` (no fastmcp).
+"""

kx_auth_cli/assert_cmd.py

@@ -0,0 +1,209 @@
+"""``kx auth assert`` — exercise the kdb+ identity-assertion handshake from a shell.
+
+plain kdb+ has no bearer over qIPC, so the container does not mint a token for it — it **asserts**
+the caller's identity: projects the validated claims into a q dict and binds them to a trusted
+service-account connection via ``.kx.auth.bind[principal]``. This command runs that handshake
+end-to-end *without* a full container, so an operator or CI step can verify a kdb+ target's
+``kx.auth`` module + permission functions behave as expected for a given principal.
+
+Two modes:
+
+- **project-only** (no ``--connect``): build and print the principal wire-dict the container *would*
+  bind. Uses the shared, fastmcp-free projection — always available, no PyKX needed.
+- **handshake** (``--connect host:port``): connect as the service account, call ``.kx.auth.bind``,
+  then read back ``.kx.auth.current[]`` / ``.kx.auth.valid[]`` (and optionally run a ``--probe``
+  query) to confirm the assertion took. The qIPC leg needs PyKX, pulled in via the optional
+  ``kx-auth-cli[qipc]`` extra so the base CLI stays light + fastmcp-free.
+
+Claims source (first found wins): ``--principal`` JSON (``-`` for stdin, ``@file`` for a file) →
+``--token`` (a JWT, decoded **unverified** — this is a local handshake exerciser, not a validator) →
+``$KX_AUTH_TOKEN`` (as a JWT). Exit codes: ``0`` ok · ``1`` error · ``2`` usage · ``4`` denied
+(a ``--probe`` query refused by the q-side permission check).
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+
+from kx_auth_core import decode_claims_unverified, project_from_claims
+
+EXIT_OK = 0
+EXIT_ERROR = 1
+EXIT_USAGE = 2
+EXIT_DENIED = 4
+
+HELP_DESCRIPTION = (
+    "Exercise the kdb+ identity-assertion handshake: project a principal into the q wire-dict and "
+    "(with --connect) bind it on a service-account connection to verify a target's kx.auth module."
+)
+HELP_EPILOG = """\
+exit codes (branch on the code, not the text):
+  0  ok        projected (project-only), or connected + bound + asserted
+  1  error     bad claims, target missing .kx.auth, connect failure, or PyKX not installed
+  2  usage     no principal/claims supplied
+  4  denied    a --probe query was refused by the q-side permission check
+
+claims source (first found wins):  --principal JSON  ->  --token JWT  ->  $KX_AUTH_TOKEN (JWT)
+  --principal '-'        read JSON claims from stdin
+  --principal @file.json read JSON claims from a file
+
+examples:
+  # project-only: see exactly what would be bound (no kdb+, no PyKX)
+  kx auth assert --principal '{"sub":"alice","scope":"kdbx.read","aud":"kx-mcp"}' --json
+
+  # full handshake against a kdb+ carrying the kx.auth module (needs kx-auth-cli[qipc])
+  kx auth assert --token "$TOK" --connect localhost:5010 \\
+    --user kxmcp --password "$SVC_PW" --probe "select from trades" --json
+"""
+
+
+def add_arguments(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--principal", help="Principal claims as JSON ('-' for stdin, '@file' for a file).")
+    parser.add_argument("--token", help="A JWT to decode (UNVERIFIED) into claims — convenience for a local handshake.")
+    parser.add_argument("--connect", metavar="HOST:PORT", help="kdb+ target to run the bind handshake against (needs the [qipc] extra).")
+    parser.add_argument("--user", default=os.environ.get("KX_AUTH_KDB_USER", ""), help="Service-account user for the qIPC login (or $KX_AUTH_KDB_USER).")
+    parser.add_argument(
+        "--password",
+        default=os.environ.get("KX_AUTH_KDB_PASSWORD"),
+        help="Service-account password (or $KX_AUTH_KDB_PASSWORD — preferred, keeps it out of shell history).",
+    )
+    parser.add_argument("--tls", action="store_true", help="Use TLS for the qIPC connection.")
+    parser.add_argument("--probe", help="Optional q expression to run after binding (a denial maps to exit 4).")
+    parser.add_argument("--timeout", type=float, default=5.0, help="qIPC connection timeout, seconds.")
+    parser.add_argument("--json", action="store_true", help="Emit a structured JSON envelope (the agent path).")
+
+
+def _resolve_claims(args: argparse.Namespace) -> tuple[dict | None, str | None]:
+    """Return (claims, error). Claims source: --principal JSON, then --token, then $KX_AUTH_TOKEN."""
+    if args.principal:
+        raw = args.principal
+        if raw == "-":
+            raw = sys.stdin.read()
+        elif raw.startswith("@"):
+            try:
+                with open(raw[1:], "r") as fh:
+                    raw = fh.read()
+            except OSError as exc:
+                return None, f"cannot read principal file: {exc}"
+        try:
+            claims = json.loads(raw)
+        except json.JSONDecodeError as exc:
+            return None, f"--principal is not valid JSON: {exc}"
+        if not isinstance(claims, dict):
+            return None, "--principal JSON must be an object of claims"
+        return claims, None
+
+    token = args.token or os.environ.get("KX_AUTH_TOKEN")
+    if token:
+        claims = decode_claims_unverified(token.strip())
+        if not claims:
+            return None, "--token did not decode to any claims (opaque or malformed JWT)"
+        return claims, None
+
+    return None, None
+
+
+def _emit(args: argparse.Namespace, *, status: str, principal=None, reason=None, **extra) -> None:
+    if args.json:
+        envelope = {"status": status}
+        if principal is not None:
+            envelope["principal"] = principal
+        if reason is not None:
+            envelope["reason"] = reason
+        envelope.update(extra)
+        print(json.dumps(envelope, indent=2, default=str))
+        return
+    if status == "ok":
+        print(f"ok — principal sub={principal.get('sub')!r}", end="")
+        if extra:
+            print(" " + " ".join(f"{k}={v}" for k, v in extra.items()))
+        else:
+            print()
+    else:
+        print(f"{status}: {reason}", file=sys.stderr)
+
+
+def _handshake(args: argparse.Namespace, wire: dict) -> int:
+    """Connect as the service account, bind the principal, confirm, optionally probe."""
+    try:
+        import pykx as kx
+    except Exception:
+        _emit(
+            args,
+            status="error",
+            principal=wire,
+            reason="PyKX is required for --connect; install the qIPC extra: pip install 'kx-auth-cli[qipc]'",
+        )
+        return EXIT_ERROR
+
+    host, _, port = args.connect.partition(":")
+    if not port.isdigit():
+        _emit(args, status="error", reason="--connect must be HOST:PORT")
+        return EXIT_USAGE
+
+    try:
+        conn = kx.SyncQConnection(
+            host=host or "localhost",
+            port=int(port),
+            username=args.user,
+            password=args.password or "",
+            timeout=args.timeout,
+            tls=args.tls,
+        )
+    except Exception as exc:
+        _emit(args, status="error", reason=f"connect failed: {exc}")
+        return EXIT_ERROR
+
+    try:
+        conn(".kx.auth.bind", wire)
+    except Exception as exc:
+        _emit(
+            args,
+            status="error",
+            principal=wire,
+            reason=f"bind failed (target missing the kx.auth module?): {exc}",
+        )
+        return EXIT_ERROR
+
+    try:
+        bound_valid = bool(conn(".kx.auth.valid[]").py())
+    except Exception as exc:
+        _emit(args, status="error", principal=wire, reason=f"could not read .kx.auth.valid[]: {exc}")
+        return EXIT_ERROR
+
+    if args.probe:
+        try:
+            conn(args.probe)
+        except Exception as exc:
+            if str(exc).strip().lower().startswith("denied"):
+                _emit(args, status="denied", principal=wire, reason=str(exc), bound=True, valid=bound_valid)
+                return EXIT_DENIED
+            _emit(args, status="error", principal=wire, reason=f"probe failed: {exc}")
+            return EXIT_ERROR
+
+    _emit(args, status="ok", principal=wire, bound=True, valid=bound_valid, probed=bool(args.probe))
+    return EXIT_OK
+
+
+def run(args: argparse.Namespace) -> int:
+    claims, err = _resolve_claims(args)
+    if err is not None:
+        _emit(args, status="error", reason=err)
+        return EXIT_ERROR
+    if claims is None:
+        print(
+            "error: no principal supplied (pass --principal JSON, --token JWT, or $KX_AUTH_TOKEN)",
+            file=sys.stderr,
+        )
+        return EXIT_USAGE
+
+    wire = project_from_claims(claims)
+
+    if not args.connect:
+        _emit(args, status="ok", principal=wire)
+        return EXIT_OK
+
+    return _handshake(args, wire)

kx_auth_cli/cache.py

@@ -0,0 +1,60 @@
+"""Client-side token cache for ``kx auth`` — a single JSON file, keyed by endpoint.
+
+``login`` writes the acquired credential here; ``exchange`` reads it as a subject-token fallback.
+The cache is a *client-side* concern (the agent/operator's own credentials), so it lives in the CLI
+rather than the shared core. Templates include product CLIs' endpoint-keyed caches and ``gh``'s
+``~/.config/gh/hosts.yml`` — a single JSON file under a dotdir, path overridable, keyed by endpoint
+so multiple deployments coexist. OS-keyring storage is a later hardening, not the v1 bar.
+
+The keys are the ``--server`` (MCP-server / resource) URLs the agent connects to. ``introspect`` is
+stateless and never touches this file.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from typing import Any, Optional
+
+
+def cache_path() -> Path:
+    """The credential file path: ``$KX_AUTH_CACHE`` if set, else ``~/.kx/credentials.json``."""
+    override = os.environ.get("KX_AUTH_CACHE")
+    if override:
+        return Path(override).expanduser()
+    return Path.home() / ".kx" / "credentials.json"
+
+
+def load() -> dict[str, Any]:
+    """The whole cache as ``{server: credential}``; ``{}`` if absent or unreadable (never raises)."""
+    path = cache_path()
+    try:
+        with path.open("r", encoding="utf-8") as fh:
+            data = json.load(fh)
+        return data if isinstance(data, dict) else {}
+    except (OSError, ValueError):
+        # Missing or corrupt cache is "no cached credentials", not a hard error.
+        return {}
+
+
+def get(server: str) -> Optional[dict[str, Any]]:
+    """The cached credential for ``server``, or ``None`` if not present."""
+    entry = load().get(server)
+    return entry if isinstance(entry, dict) else None
+
+
+def save(server: str, credential: dict[str, Any]) -> Path:
+    """Merge ``credential`` under ``server`` and persist. Dir ``0700``, file ``0600``."""
+    path = cache_path()
+    path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
+    data = load()
+    data[server] = credential
+    # Write via a temp file then replace, so a crash mid-write can't truncate the cache. The temp
+    # file is created 0600 at open (not chmod'd after), so tokens never sit on disk umask-readable.
+    tmp = path.with_suffix(path.suffix + ".tmp")
+    fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+    with os.fdopen(fd, "w", encoding="utf-8") as fh:
+        json.dump(data, fh, indent=2, default=str)
+    os.replace(tmp, path)
+    return path

kx_auth_cli/cli.py

@@ -0,0 +1,85 @@
+"""The ``kx`` umbrella entry point (stdlib argparse).
+
+The ``auth`` group exposes ``introspect``, ``login``, ``exchange``, and ``assert``; the umbrella
+leaves room for future top-level ``kx`` commands. Each leaf command sets ``func``; ``main`` dispatches
+and propagates its int return as the process exit code. Argparse's own usage errors exit ``2`` — the
+code the spec reserves for that.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from typing import Optional, Sequence
+
+from . import assert_cmd, exchange, introspect, login
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="kx",
+        description="kx control-plane CLI. Agent-first: every command supports --json and a stable "
+        "exit-code contract (0 ok · 1 error · 2 usage · 3 auth-required · 4 denied).",
+    )
+    groups = parser.add_subparsers(dest="group", metavar="<group>")
+
+    auth = groups.add_parser(
+        "auth",
+        help="Authentication flows (introspect / login / exchange / assert).",
+        description="Authentication flows for the kx-mcp control plane: `introspect`, `login`, "
+        "`exchange`, `assert`.",
+    )
+    auth_commands = auth.add_subparsers(dest="auth_command", metavar="<command>")
+
+    introspect_parser = auth_commands.add_parser(
+        "introspect",
+        help="Validate a bearer token against the configured KX_MCP_AUTH keys.",
+        description=introspect.HELP_DESCRIPTION,
+        epilog=introspect.HELP_EPILOG,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    introspect.add_arguments(introspect_parser)
+    introspect_parser.set_defaults(func=introspect.run)
+
+    login_parser = auth_commands.add_parser(
+        "login",
+        help="Acquire a bearer for an MCP server via server-advertised device-code OAuth.",
+        description=login.HELP_DESCRIPTION,
+        epilog=login.HELP_EPILOG,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    login.add_arguments(login_parser)
+    login_parser.set_defaults(func=login.run)
+
+    exchange_parser = auth_commands.add_parser(
+        "exchange",
+        help="Swap a subject token for a backend-scoped credential (RFC 8693).",
+        description=exchange.HELP_DESCRIPTION,
+        epilog=exchange.HELP_EPILOG,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    exchange.add_arguments(exchange_parser)
+    exchange_parser.set_defaults(func=exchange.run)
+
+    assert_parser = auth_commands.add_parser(
+        "assert",
+        help="Exercise the kdb+ identity-assertion handshake (project a principal; optionally bind it).",
+        description=assert_cmd.HELP_DESCRIPTION,
+        epilog=assert_cmd.HELP_EPILOG,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    assert_cmd.add_arguments(assert_parser)
+    assert_parser.set_defaults(func=assert_cmd.run)
+
+    return parser
+
+
+def main(argv: Optional[Sequence[str]] = None) -> None:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    func = getattr(args, "func", None)
+    if func is None:
+        # `kx` or `kx auth` with no leaf command — show help and exit as a usage error.
+        parser.print_help(sys.stderr)
+        sys.exit(2)
+    sys.exit(func(args))

kx_auth_cli/discovery.py

@@ -0,0 +1,231 @@
+"""OAuth discovery + device-code primitives for ``kx auth login`` (stdlib + httpx, fastmcp-free).
+
+The mediation model is fixed by the MCP spec and `mcp-container-design.md`: the agent points at an
+endpoint it already knows — **the MCP server (the container / resource server)** — and that endpoint
+advertises its authorization server. So the flow is:
+
+1. **RFC 9728** — GET the protected-resource metadata → ``authorization_servers``. The path-aware
+   form the RFC specifies is tried first (``/.well-known/oauth-protected-resource/<path>`` for a
+   server mounted under a path), falling back to the origin-root form.
+2. **RFC 8414 / OIDC** — GET the AS's ``.well-known`` metadata → device + token (+ registration)
+   endpoints.
+3. **RFC 7591 (optional)** — Dynamic Client Registration when the AS advertises a
+   ``registration_endpoint`` and no ``--client-id`` was given (the spec's "user never hand-configures
+   a client-id" ideal).
+4. **RFC 8628** — device-authorization request, then poll the token endpoint.
+
+After discovery the client talks **directly to the authorization server** — never back through the
+container (the spec forbids the resource server passing the token through). Errors carry a ``reason``
+code so the command layer can map them to the stable exit-code contract.
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import Any, Callable, Optional
+from urllib.parse import urljoin, urlsplit, urlunsplit
+
+import httpx
+
+
+class DiscoveryError(Exception):
+    """RFC 9728/8414 discovery failed (no metadata, no AS, no required endpoint)."""
+
+
+class DeviceAuthError(Exception):
+    """The device-code flow failed. ``reason`` is the OAuth error code (or a synthetic one).
+
+    Reasons the command layer maps to exit codes: ``access_denied`` → 4 (denied);
+    ``expired_token`` / ``timeout`` → 3 (auth-required); everything else → 1 (error).
+    """
+
+    def __init__(self, message: str, *, reason: str = "error") -> None:
+        super().__init__(message)
+        self.reason = reason
+
+
+@dataclass
+class AuthMetadata:
+    """The authorization-server endpoints `login` needs, resolved from discovery."""
+
+    authorization_server: str
+    token_endpoint: str
+    device_authorization_endpoint: Optional[str] = None
+    registration_endpoint: Optional[str] = None
+
+
+def _prm_candidates(server: str) -> list[str]:
+    """RFC 9728 well-known URLs for a resource: the path-aware form first, then the origin root."""
+    parts = urlsplit(server)
+    origin = urlunsplit((parts.scheme, parts.netloc, "", "", ""))
+    root = origin + "/.well-known/oauth-protected-resource"
+    path = parts.path.strip("/")
+    return [f"{root}/{path}", root] if path else [root]
+
+
+def _get_json(client: httpx.Client, url: str) -> dict[str, Any]:
+    resp = client.get(url, headers={"Accept": "application/json"})
+    resp.raise_for_status()
+    try:
+        body = resp.json()
+    except ValueError as exc:
+        raise DiscoveryError(f"metadata at {url} was not valid JSON") from exc
+    if not isinstance(body, dict):
+        raise DiscoveryError(f"metadata at {url} was not a JSON object")
+    return body
+
+
+def discover(server: str, *, client: httpx.Client) -> AuthMetadata:
+    """Resolve the AS endpoints for an MCP-server (resource) URL via RFC 9728 → RFC 8414/OIDC."""
+    prm: Optional[dict[str, Any]] = None
+    last_prm_exc: Optional[Exception] = None
+    for prm_url in _prm_candidates(server):
+        try:
+            candidate = _get_json(client, prm_url)
+        except (httpx.HTTPError, DiscoveryError) as exc:
+            last_prm_exc = exc
+            continue
+        if candidate.get("authorization_servers"):
+            prm = candidate
+            break
+        last_prm_exc = DiscoveryError(
+            f"{prm_url} advertised no authorization_servers — cannot discover an IdP"
+        )
+    if prm is None:
+        raise DiscoveryError(
+            f"could not discover protected-resource metadata for {server}: {last_prm_exc}"
+        )
+    as_url = prm["authorization_servers"][0]
+
+    # RFC 8414 and OIDC differ in well-known path; try both against the AS origin + issuer path.
+    candidates = [
+        urljoin(as_url.rstrip("/") + "/", ".well-known/oauth-authorization-server"),
+        urljoin(as_url.rstrip("/") + "/", ".well-known/openid-configuration"),
+    ]
+    last_exc: Optional[Exception] = None
+    for url in candidates:
+        try:
+            meta = _get_json(client, url)
+        except (httpx.HTTPError, DiscoveryError) as exc:
+            last_exc = exc
+            continue
+        token_endpoint = meta.get("token_endpoint")
+        if not token_endpoint:
+            last_exc = DiscoveryError(f"{url} advertised no token_endpoint")
+            continue
+        return AuthMetadata(
+            authorization_server=as_url,
+            token_endpoint=token_endpoint,
+            device_authorization_endpoint=meta.get("device_authorization_endpoint"),
+            registration_endpoint=meta.get("registration_endpoint"),
+        )
+    raise DiscoveryError(
+        f"could not fetch authorization-server metadata from {as_url}: {last_exc}"
+    )
+
+
+def register_client(
+    meta: AuthMetadata, *, client: httpx.Client, client_name: str, scopes: Optional[list[str]] = None
+) -> str:
+    """RFC 7591 Dynamic Client Registration for the device-code flow → the issued ``client_id``."""
+    if not meta.registration_endpoint:
+        raise DiscoveryError(
+            "the authorization server advertises no registration_endpoint — pass --client-id "
+            "(or $KX_AUTH_CLIENT_ID) instead"
+        )
+    payload: dict[str, Any] = {
+        "client_name": client_name,
+        "grant_types": ["urn:ietf:params:oauth:grant-type:device_code"],
+        "token_endpoint_auth_method": "none",  # public client — device-code, no secret
+    }
+    if scopes:
+        payload["scope"] = " ".join(scopes)
+    try:
+        resp = client.post(meta.registration_endpoint, json=payload)
+        resp.raise_for_status()
+        body = resp.json()
+    except (httpx.HTTPError, ValueError) as exc:  # ValueError: a non-JSON body
+        raise DiscoveryError(f"dynamic client registration failed: {exc}") from exc
+    client_id = body.get("client_id")
+    if not client_id:
+        raise DiscoveryError("registration endpoint returned no client_id")
+    return client_id
+
+
+def start_device_authorization(
+    meta: AuthMetadata, client_id: str, *, client: httpx.Client, scopes: Optional[list[str]] = None
+) -> dict[str, Any]:
+    """RFC 8628 device-authorization request → the device/user codes + polling parameters."""
+    if not meta.device_authorization_endpoint:
+        raise DeviceAuthError(
+            "the authorization server advertises no device_authorization_endpoint — "
+            "device-code login is unavailable",
+            reason="device_flow_unsupported",
+        )
+    data = {"client_id": client_id}
+    if scopes:
+        data["scope"] = " ".join(scopes)
+    try:
+        resp = client.post(meta.device_authorization_endpoint, data=data)
+        resp.raise_for_status()
+        body = resp.json()
+    except (httpx.HTTPError, ValueError) as exc:  # ValueError: a non-JSON body
+        raise DeviceAuthError(f"device-authorization request failed: {exc}") from exc
+    if not body.get("device_code") or not body.get("user_code"):
+        raise DeviceAuthError("device-authorization response missing device_code/user_code")
+    return body
+
+
+def poll_for_token(
+    meta: AuthMetadata,
+    client_id: str,
+    device_code: str,
+    *,
+    client: httpx.Client,
+    interval: int = 5,
+    expires_in: int = 600,
+    sleep: Callable[[float], None] = time.sleep,
+    monotonic: Callable[[], float] = time.monotonic,
+) -> dict[str, Any]:
+    """RFC 8628 polling: hit the token endpoint until success, denial, or expiry.
+
+    ``sleep`` / ``monotonic`` are injectable so tests drive the loop without real waits.
+    """
+    deadline = monotonic() + expires_in
+    wait = max(1, interval)
+    while True:
+        resp = client.post(
+            meta.token_endpoint,
+            data={
+                "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+                "device_code": device_code,
+                "client_id": client_id,
+            },
+        )
+        try:
+            body: dict[str, Any] = resp.json() if resp.content else {}
+        except ValueError:
+            # A non-JSON body (e.g. a proxy's HTML error page) → fall through to the status-code
+            # error below rather than crashing out of the exit-code contract.
+            body = {}
+        if resp.status_code == 200 and body.get("access_token"):
+            return body
+
+        error = body.get("error", "")
+        if error == "authorization_pending":
+            pass
+        elif error == "slow_down":
+            wait += 5
+        elif error == "access_denied":
+            raise DeviceAuthError("the user denied the authorization request", reason="access_denied")
+        elif error == "expired_token":
+            raise DeviceAuthError("the device code expired before approval", reason="expired_token")
+        else:
+            raise DeviceAuthError(
+                f"token endpoint returned an unexpected error: {error or resp.status_code}"
+            )
+
+        if monotonic() >= deadline:
+            raise DeviceAuthError("timed out waiting for device-code approval", reason="timeout")
+        sleep(wait)