agent-credential-broker/agent_credential_broker/server.py

"""Unix-domain socket daemon. No TCP, ever.

Socket: $XDG_RUNTIME_DIR/agent-broker/broker.sock (0600,
dir 0700). Every request is one line of JSON:

    {"method": "...", "params": {...}, "id": 1}

Response:

    {"ok": true, "result": {...}}
    {"ok": false, "error": {"code": "...", "message": "..."}}

The daemon additionally verifies the connecting peer UID via
SO_PEERCRED -- only the owning user may talk to it.
"""

from __future__ import annotations

import json
import os
import socket
import socketserver
import struct
import threading

from .config import Config, socket_path
from .core import BrokerCore
from .errors import BrokerError, internal
from .audit import audit

_MAX_LINE = 1024 * 1024


def _peer_uid(conn: socket.socket) -> int:
    creds = conn.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, struct.calcsize("3i"))
    _pid, uid, _gid = struct.unpack("3i", creds)
    return uid


class _Handler(socketserver.StreamRequestHandler):
    def handle(self):
        try:
            if _peer_uid(self.request) != os.getuid():
                self.wfile.write(b'{"ok":false,"error":{"code":"DENIED","message":"foreign uid"}}\n')
                return
        except OSError:
            return
        for raw in self.rfile:
            if len(raw) > _MAX_LINE:
                break
            try:
                req = json.loads(raw)
            except json.JSONDecodeError:
                self._send({"ok": False, "error": {"code": "DENIED_INVALID_ARGUMENT", "message": "bad json"}})
                continue
            self._send(self.server.dispatch(req))

    def _send(self, obj):
        try:
            self.wfile.write(json.dumps(obj).encode() + b"\n")
        except OSError:
            pass


class _Server(socketserver.ThreadingMixIn, socketserver.UnixStreamServer):
    daemon_threads = True

    def __init__(self, path: str, core: BrokerCore):
        self.core = core
        super().__init__(path, _Handler)

    def dispatch(self, req: dict) -> dict:
        if not isinstance(req, dict):
            return _err(internal("bad request"))
        method = req.get("method")
        params = req.get("params") or {}
        if not isinstance(params, dict):
            return _err(BrokerError("DENIED_INVALID_ARGUMENT", "params must be an object"))
        try:
            result = self._route(method, params)
            return {"ok": True, "result": result}
        except BrokerError as e:
            out = {"ok": False, "error": e.to_dict()}
            if getattr(e, "ticket", None):
                out["ticket"] = e.ticket
            return out
        except Exception:
            return _err(internal("unexpected daemon error"))

    def _route(self, method, params):
        c = self.core
        routes = {
            "status": lambda: c.status(),
            "unlock": lambda: c.unlock(params.get("password") or _bad("password required")),
            "lock": lambda: c.lock(),
            "logout": lambda: c.logout(),
            "confirm": lambda: c.confirm(params.get("ticket") or _bad("ticket required")),
            "list_refs": lambda: c.list_refs(),
            "find_for_url": lambda: c.find_for_url(params.get("url") or _bad("url required")),
            # internal path used by the agent-browser credential.read plugin
            "resolve": lambda: c.resolve_for_plugin(
                params.get("ref") or _bad("ref required"), params.get("url")),
            "http_request": lambda: c.run("http", params),
            "ssh_exec": lambda: c.run("ssh", params),
            "git_op": lambda: c.run("git", params),
            "cli_run": lambda: c.run("cli-profile", params),
            "browser_login": lambda: c.run("browser", {**params, "op": "browser_login"}),
            "browser_fill_credential": lambda: c.run("browser", {**params, "op": "browser_fill_credential"}),
        }
        fn = routes.get(method)
        if fn is None:
            raise BrokerError("DENIED_INVALID_ARGUMENT", f"unknown method {method!r}")
        return fn()


def _bad(msg):
    raise BrokerError("DENIED_INVALID_ARGUMENT", msg)


def _err(e: BrokerError) -> dict:
    return {"ok": False, "error": e.to_dict()}


def call(params_method: str, params: dict | None = None, *, timeout=60) -> dict:
    """Client-side call into the daemon. Used by the CLI, the MCP server
    and the agent-browser plugin."""
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    sock.settimeout(timeout)
    try:
        sock.connect(str(socket_path()))
        sock.sendall(json.dumps({"method": params_method, "params": params or {}}).encode() + b"\n")
        fh = sock.makefile("r")
        line = fh.readline(_MAX_LINE)
        if not line:
            raise BrokerError("DAEMON_DOWN", "daemon closed connection")
        resp = json.loads(line)
    except (ConnectionRefusedError, FileNotFoundError, socket.timeout) as e:
        raise BrokerError("DAEMON_DOWN", f"cannot reach broker daemon ({e.__class__.__name__}); is it running?")
    finally:
        sock.close()
    if not resp.get("ok"):
        err = resp.get("error") or {}
        e = BrokerError(err.get("code", "INTERNAL_ERROR"), err.get("message", "unknown"))
        if resp.get("ticket"):
            e.ticket = resp["ticket"]
        raise e
    return resp.get("result")


def serve_forever(config: Config | None = None):
    config = config or Config.load()
    path = socket_path()
    path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
    try:
        path.unlink()
    except FileNotFoundError:
        pass
    core = BrokerCore(config)
    server = _Server(str(path), core)
    os.chmod(path, 0o600)
    audit("daemon_start", result="ok", detail=str(path))
    try:
        server.serve_forever(poll_interval=0.5)
    finally:
        core.lock()
        try:
            path.unlink()
        except OSError:
            pass
        audit("daemon_stop", result="ok")