UmbraAudio/infra/remote_credential_broker/keychain.py

from __future__ import annotations

import ctypes
import re
import sys
import threading
import time
from ctypes import POINTER, c_int32, c_uint32, c_void_p, create_string_buffer
from typing import Final


_PROFILE_RE: Final = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]{0,127}$")
_SERVICE_DEFAULT: Final = "com.umbratrans.remote-broker"
_ERR_SEC_ITEM_NOT_FOUND: Final = -25300


class CredentialError(RuntimeError):
    """Base exception for credential-store failures."""


class CredentialNotFound(CredentialError):
    """Raised when a requested Keychain item does not exist."""


class UnsupportedPlatformError(CredentialError):
    """Raised when the broker is not running on macOS."""


def validate_profile(profile: str) -> str:
    """Validate a profile name before using it as a Keychain account name."""
    if not isinstance(profile, str) or not _PROFILE_RE.fullmatch(profile):
        raise CredentialError(
            "Credential profile must contain only letters, digits, '.', '_' and '-' "
            "and must be 1-128 characters"
        )
    return profile


class MacKeychain:
    """Small ctypes wrapper around the legacy macOS generic-password API.

    Password bytes are passed directly to the Security framework and are never
    placed in a subprocess argument, environment variable, log line, or MCP
    response. The API is deprecated by Apple but remains available on the
    supported macOS hosts and avoids adding a Python dependency.
    """

    def __init__(self, service: str = _SERVICE_DEFAULT) -> None:
        self.service = service
        self._lock = threading.RLock()
        self._lib: ctypes.CDLL | None = None
        self._cf: ctypes.CDLL | None = None

    def _load(self) -> tuple[ctypes.CDLL, ctypes.CDLL]:
        if sys.platform != "darwin":
            raise UnsupportedPlatformError(
                "The macOS Keychain backend requires macOS; use a native secret store on other hosts"
            )
        if self._lib is None:
            lib = ctypes.CDLL("/System/Library/Frameworks/Security.framework/Security")
            lib.SecKeychainAddGenericPassword.argtypes = [
                c_void_p,
                c_uint32,
                c_void_p,
                c_uint32,
                c_void_p,
                c_uint32,
                c_void_p,
                POINTER(c_void_p),
            ]
            lib.SecKeychainAddGenericPassword.restype = c_int32
            lib.SecKeychainFindGenericPassword.argtypes = [
                c_void_p,
                c_uint32,
                c_void_p,
                c_uint32,
                c_void_p,
                POINTER(c_uint32),
                POINTER(c_void_p),
                POINTER(c_void_p),
            ]
            lib.SecKeychainFindGenericPassword.restype = c_int32
            lib.SecKeychainItemModifyAttributesAndData.argtypes = [
                c_void_p,
                c_void_p,
                c_uint32,
                c_void_p,
            ]
            lib.SecKeychainItemModifyAttributesAndData.restype = c_int32
            lib.SecKeychainItemFreeContent.argtypes = [c_void_p, c_void_p]
            lib.SecKeychainItemFreeContent.restype = c_int32
            lib.SecKeychainItemDelete.argtypes = [c_void_p]
            lib.SecKeychainItemDelete.restype = c_int32
            self._lib = lib
        if self._cf is None:
            cf = ctypes.CDLL("/System/Library/Frameworks/CoreFoundation.framework/CoreFoundation")
            cf.CFRelease.argtypes = [c_void_p]
            cf.CFRelease.restype = None
            self._cf = cf
        return self._lib, self._cf

    @staticmethod
    def _error(status: int, operation: str) -> CredentialError:
        return CredentialError(f"Keychain {operation} failed with OSStatus {status}")

    def _find(
        self,
        profile: str,
        *,
        include_password: bool,
    ) -> tuple[int, bytes | None, c_void_p | None]:
        lib, _ = self._load()
        profile = validate_profile(profile)
        service_bytes = self.service.encode("utf-8")
        account_bytes = profile.encode("utf-8")
        service_buf = create_string_buffer(service_bytes)
        account_buf = create_string_buffer(account_bytes)
        length = c_uint32(0)
        data = c_void_p()
        item = c_void_p()
        status = lib.SecKeychainFindGenericPassword(
            None,
            len(service_bytes),
            service_buf,
            len(account_bytes),
            account_buf,
            ctypes.byref(length) if include_password else None,
            ctypes.byref(data) if include_password else None,
            ctypes.byref(item),
        )
        if status != 0:
            return status, None, None
        if not include_password:
            return status, None, item
        try:
            value = ctypes.string_at(data, length.value) if data.value else b""
            return status, value, item
        finally:
            # The pointer is owned by the Security framework.
            lib.SecKeychainItemFreeContent(None, data)

    def _release(self, item: c_void_p | None) -> None:
        if item is None or not item.value:
            return
        _, cf = self._load()
        cf.CFRelease(item)

    def set(self, profile: str, secret: str) -> None:
        """Store or replace a secret in the default login Keychain."""
        if not isinstance(secret, str) or not secret:
            raise CredentialError("Credential secret must not be empty")
        with self._lock:
            lib, _ = self._load()
            profile = validate_profile(profile)
            status, _, item = self._find(profile, include_password=False)
            if status == 0 and item is not None:
                secret_bytes = secret.encode("utf-8")
                secret_buf = create_string_buffer(secret_bytes)
                try:
                    update_status = lib.SecKeychainItemModifyAttributesAndData(
                        item,
                        None,
                        len(secret_bytes),
                        secret_buf,
                    )
                finally:
                    self._release(item)
                if update_status != 0:
                    raise self._error(update_status, "update")
                return
            if status != _ERR_SEC_ITEM_NOT_FOUND:
                raise self._error(status, "lookup")

            service_bytes = self.service.encode("utf-8")
            account_bytes = profile.encode("utf-8")
            secret_bytes = secret.encode("utf-8")
            service_buf = create_string_buffer(service_bytes)
            account_buf = create_string_buffer(account_bytes)
            secret_buf = create_string_buffer(secret_bytes)
            new_item = c_void_p()
            add_status = lib.SecKeychainAddGenericPassword(
                None,
                len(service_bytes),
                service_buf,
                len(account_bytes),
                account_buf,
                len(secret_bytes),
                secret_buf,
                ctypes.byref(new_item),
            )
            self._release(new_item)
            if add_status != 0:
                raise self._error(add_status, "add")

    def get(self, profile: str) -> str:
        """Retrieve a secret for internal broker use only."""
        with self._lock:
            status, value, item = self._find(profile, include_password=True)
            self._release(item)
            if status == _ERR_SEC_ITEM_NOT_FOUND:
                raise CredentialNotFound(f"No stored credential for profile {profile!r}")
            if status != 0:
                raise self._error(status, "lookup")
            if value is None:
                raise CredentialError("Keychain returned an empty credential")
            return value.decode("utf-8")

    def delete(self, profile: str) -> None:
        """Delete a stored secret; missing items are treated as already absent."""
        with self._lock:
            lib, _ = self._load()
            status, _, item = self._find(profile, include_password=False)
            if status == _ERR_SEC_ITEM_NOT_FOUND:
                return
            if status != 0:
                raise self._error(status, "lookup")
            if item is None:
                return
            try:
                delete_status = lib.SecKeychainItemDelete(item)
            finally:
                self._release(item)
            if delete_status != 0:
                raise self._error(delete_status, "delete")

    def exists(self, profile: str) -> bool:
        """Return whether a profile exists without retrieving its secret."""
        with self._lock:
            status, _, item = self._find(profile, include_password=False)
            self._release(item)
            if status == 0:
                return True
            if status == _ERR_SEC_ITEM_NOT_FOUND:
                return False
            raise self._error(status, "lookup")


class CredentialVault:
    """Persistent Keychain storage plus explicitly temporary in-memory profiles."""

    def __init__(self, persistent: MacKeychain | None = None) -> None:
        self.persistent = persistent or MacKeychain()
        self.service = self.persistent.service
        self._lock = threading.RLock()
        self._temporary: dict[str, tuple[str, float]] = {}
        self._timers: dict[str, threading.Timer] = {}

    def _expire(self, profile: str) -> None:
        with self._lock:
            self._temporary.pop(profile, None)
            self._timers.pop(profile, None)

    def set(self, profile: str, secret: str, ttl_seconds: int | None = None) -> None:
        profile = validate_profile(profile)
        if not isinstance(secret, str) or not secret:
            raise CredentialError("Credential secret must not be empty")
        with self._lock:
            old_timer = self._timers.pop(profile, None)
            if old_timer is not None:
                old_timer.cancel()
            if ttl_seconds is None:
                self._temporary.pop(profile, None)
                self.persistent.set(profile, secret)
                return
            if not 1 <= int(ttl_seconds) <= 3600:
                raise CredentialError("ttl_seconds must be between 1 and 3600")
            expires_at = time.monotonic() + int(ttl_seconds)
            self._temporary[profile] = (secret, expires_at)
            timer = threading.Timer(int(ttl_seconds), self._expire, args=(profile,))
            timer.daemon = True
            self._timers[profile] = timer
            timer.start()

    def get(self, profile: str) -> str:
        profile = validate_profile(profile)
        with self._lock:
            temporary = self._temporary.get(profile)
            if temporary is not None:
                secret, expires_at = temporary
                if expires_at > time.monotonic():
                    return secret
                self._expire(profile)
            return self.persistent.get(profile)

    def delete(self, profile: str) -> None:
        profile = validate_profile(profile)
        with self._lock:
            self._expire(profile)
            self.persistent.delete(profile)

    def exists(self, profile: str) -> bool:
        profile = validate_profile(profile)
        with self._lock:
            temporary = self._temporary.get(profile)
            if temporary is not None:
                if temporary[1] > time.monotonic():
                    return True
                self._expire(profile)
            return self.persistent.exists(profile)