agent-credential-broker/agent_credential_broker/core.py

"""Broker core: vault state, ref resolution, confirmations, dispatch.

This object lives inside the daemon process and is the ONLY place where
decrypted credential material exists -- in RAM, under a LOCKED default,
auto-expiring after ``unlock_ttl_seconds``.
"""

from __future__ import annotations

import secrets as _secrets
import threading
import time

from . import errors, policy as pol
from .audit import audit
from .backends import make_backend
from .config import Config
from .items import wipe_item


class VaultState:
    def __init__(self):
        self.items: list[dict] = []
        self.ref_index: dict[str, list[int]] = {}
        self.unlocked_at: float | None = None
        self.session_approvals: set = set()
        self.tickets: dict[str, dict] = {}

    def clear(self):
        for item in self.items:
            for k, v in item.items():
                item[k] = None
        self.items = []
        self.ref_index = {}
        self.unlocked_at = None
        self.session_approvals = set()
        self.tickets = {}

    @property
    def unlocked(self) -> bool:
        return self.unlocked_at is not None


class BrokerCore:
    def __init__(self, config: Config):
        self.config = config
        self.backend = make_backend(config)
        self.state = VaultState()
        self._lock = threading.RLock()

    # -- lifecycle -----------------------------------------------------

    def status(self) -> dict:
        with self._lock:
            self._expire_if_needed()
            return {
                "locked": not self.state.unlocked,
                "configured": self._configured(),
                "refs": len(self.state.ref_index) if self.state.unlocked else None,
                "backend": self.backend.name,
                "socket": "unix",
            }

    def _configured(self) -> bool:
        return self.backend.configured()

    def _expire_if_needed(self):
        st = self.state
        if st.unlocked and (time.time() - st.unlocked_at) > self.config.unlock_ttl_seconds:
            st.clear()
            self.backend.lock()
            audit("auto_lock", result="ok", detail="ttl expired")

    def unlock(self, password: str) -> dict:
        if not self._configured():
            raise errors.unconfigured()
        try:
            self.backend.unlock(password)
        finally:
            password = None
        items = self.backend.list_metadata()
        with self._lock:
            self.state.clear()
            self.state.items = items
            for i, item in enumerate(items):
                p = pol.policy_for(item)
                if p:
                    self.state.ref_index.setdefault(p.ref, []).append(i)
            self.state.unlocked_at = time.time()
            dupes = [r for r, idx in self.state.ref_index.items() if len(idx) > 1]
            audit("unlock", result="ok", detail=f"items={len(items)} refs={len(self.state.ref_index)}")
            return {
                "unlocked": True,
                "items": len(items),
                "refs": len(self.state.ref_index),
                "duplicate_refs": sorted(dupes),
            }

    def lock(self) -> dict:
        with self._lock:
            self.state.clear()
            self.backend.lock()
        audit("lock", result="ok")
        return {"locked": True}

    def logout(self) -> dict:
        with self._lock:
            self.state.clear()
            self.backend.logout()
        audit("logout", result="ok")
        return {"locked": True, "logged_out": True}

    # -- refs -----------------------------------------------------------

    def _require_unlocked(self):
        self._expire_if_needed()
        if not self.state.unlocked:
            raise errors.locked()

    def _resolve_index(self, ref: str) -> int:
        """Exact-match ref -> item index. 0 -> NOT_FOUND, >1 -> AMBIGUOUS_REF."""
        idxs = self.state.ref_index.get(ref, [])
        if not idxs:
            raise errors.not_found(ref)
        if len(idxs) > 1:
            raise errors.ambiguous(ref, len(idxs))
        return idxs[0]

    def resolve_ref(self, ref: str) -> tuple[dict, pol.CredentialPolicy]:
        """Internal resolve: returns (decrypted item, policy). Callers in
        the process only -- the socket layer restricts who may receive
        the secret material this yields."""
        self._require_unlocked()
        if not isinstance(ref, str) or not ref or len(ref) > 200:
            raise errors.invalid_argument("bad ref")
        item = self.state.items[self._resolve_index(ref)]
        p = pol.policy_for(item)
        if p is None:
            raise errors.not_found(ref)
        return item, p

    def list_refs(self) -> list[dict]:
        self._require_unlocked()
        out = []
        seen = set()
        for ref, idxs in self.state.ref_index.items():
            for i in idxs:
                p = pol.policy_for(self.state.items[i])
                meta = p.safe_metadata()
                if len(idxs) > 1:
                    meta["duplicate"] = True
                if ref not in seen:
                    out.append(meta)
                    seen.add(ref)
        return sorted(out, key=lambda m: m["ref"])

    def find_for_url(self, url: str) -> list[str]:
        self._require_unlocked()
        refs = []
        for ref in self.state.ref_index:
            try:
                item = self.state.items[self.state.ref_index[ref][0]]
                p = pol.policy_for(item)
                if p is None:
                    continue
                pol.require_url_allowed(p, url, require_https=False)
                refs.append(ref)
            except errors.BrokerError:
                continue
        return sorted(set(refs))

    # -- confirmation ----------------------------------------------------

    def _confirm(self, policy, adapter: str, target: str, ticket=None):
        level = policy.confirmation
        if level == "DENY":
            raise errors.denied(f"credential {policy.ref!r} is DENY for confirmation")
        if level == "AUTO":
            return
        if level == "SESSION" and (policy.ref, adapter) in self.state.session_approvals:
            return
        if ticket:
            tk = self.state.tickets.pop(ticket, None)
            if (tk and tk["approved"] and time.time() < tk["expires"]
                    and tk["ref"] == policy.ref and tk["adapter"] == adapter):
                if level == "SESSION":
                    self.state.session_approvals.add((policy.ref, adapter))
                return
            raise errors.denied("confirmation ticket invalid or expired")
        ticket_id = _secrets.token_urlsafe(24)
        self.state.tickets[ticket_id] = {
            "ref": policy.ref, "adapter": adapter, "target": target,
            "approved": False, "expires": time.time() + self.config.confirmation_ttl_seconds,
        }
        audit("confirmation_pending", ref=policy.ref, adapter=adapter, target=target)
        raise errors.confirmation_required(ticket_id)

    def confirm(self, ticket: str) -> dict:
        """Approve a pending ticket. Invoked interactively by the human."""
        tk = self.state.tickets.get(ticket)
        if not tk or time.time() > tk["expires"]:
            raise errors.denied("unknown or expired ticket")
        tk["approved"] = True
        tk["expires"] = time.time() + self.config.confirmation_ttl_seconds
        audit("confirmed", ref=tk["ref"], adapter=tk["adapter"], target=tk["target"])
        return {"approved": True, "ref": tk["ref"], "adapter": tk["adapter"], "target": tk["target"]}

    # -- adapter dispatch --------------------------------------------------

    def run(self, adapter: str, params: dict) -> dict:
        """Central guarded dispatch used by every adapter."""
        self._require_unlocked()
        ref = params.get("credential_ref") or params.get("ref")
        ticket = params.get("ticket")
        meta, p = self.resolve_ref(ref)
        pol.require_adapter(p, adapter)
        target = params.get("url") or params.get("host") or params.get("repository") or params.get("profile") or ""
        from .adapters import registry
        registry.validate(adapter, meta, p, params, self.config)
        self._confirm(p, adapter, target, ticket=ticket)
        item = self.backend.get_item(meta["id"])
        try:
            result = registry.run(adapter, item, p, params, self.config)
            audit(adapter, ref=p.ref, adapter=adapter, target=target, result="ok")
            return result
        except errors.BrokerError as e:
            audit(adapter, ref=p.ref, adapter=adapter, target=target, result="error", error_code=e.code)
            raise
        except Exception:
            audit(adapter, ref=p.ref, adapter=adapter, target=target, result="error", error_code="INTERNAL_ERROR")
            raise errors.internal("adapter failed")
        finally:
            wipe_item(item)

    def resolve_for_plugin(self, ref: str, url: str | None) -> dict:
        """Socket-exposed resolve for the agent-browser credential.read
        plugin. Only usable for refs whose policy allows the browser
        adapter; origin is validated when a URL is provided."""
        meta, p = self.resolve_ref(ref)
        pol.require_adapter(p, "browser")
        if url:
            pol.require_url_allowed(p, url)
        audit("resolve", ref=p.ref, adapter="browser", target=url or p.uris[:1])
        item = self.backend.get_item(meta["id"])
        try:
            secrets = {
                "username": item.get("username"),
                "password": item.get("password"),
                "url": url or (p.uris[0] if p.uris else None),
                "usernameSelector": p.username_selector,
                "passwordSelector": p.password_selector,
                "submitSelector": p.submit_selector,
            }
            return {k: v for k, v in secrets.items() if v}
        finally:
            wipe_item(item)