agent-identity-provisioning/bw_helper.py

#!/usr/bin/env python3
"""
bw_helper.py -- Bitwarden CLI wrapper for agent identity provisioning.

Provides a clean Python interface to the `bw` CLI for:
- Password generation
- Item creation/retrieval/updating in collections
- TOTP code generation
- Session management

All BW commands run via subprocess. The BW session is established once
and reused across calls.

SAFETY RULES:
- This module NEVER deletes items. There is no delete_item method.
- create_item() refuses to create duplicates.
- update_item() modifies an existing item in place by ID.
- get_item_id() is the canonical way to resolve an item -- it uses the
  BW search API and raises on ambiguity (multiple matches).
"""

import json
import os
import subprocess
import sys
import tempfile
import time
from typing import Optional


class BitwardenHelper:
    """Wrapper around the Bitwarden CLI for credential management."""

    def __init__(self, client_id: str, client_secret: str, password: str,
                 totp_secret: str = "", server_url: str = ""):
        self.client_id = client_id
        self.client_secret = client_secret
        self.password = password
        self.totp_secret = totp_secret
        self.server_url = server_url
        self.session: Optional[str] = None
        self._last_sync: float = 0.0

    def _run_bw(self, args: list[str], capture: bool = True) -> str:
        """Run a bw CLI command with the active session."""
        env = os.environ.copy()
        if self.session:
            env["BW_SESSION"] = self.session
        result = subprocess.run(
            ["bw"] + args,
            capture_output=capture,
            text=True,
            env=env,
        )
        if result.returncode != 0:
            raise RuntimeError(
                f"bw {' '.join(args)} failed: {result.stderr.strip()}"
            )
        return result.stdout.strip() if capture else ""

    def _encode(self, data: dict) -> str:
        """Encode a dict to base64 for bw CLI input."""
        encoded = json.dumps(data)
        result = subprocess.run(
            ["bw", "encode"],
            input=encoded,
            capture_output=True,
            text=True,
        )
        if result.returncode != 0:
            raise RuntimeError(f"bw encode failed: {result.stderr.strip()}")
        return result.stdout.strip()

    def login(self) -> None:
        """Authenticate via API key and unlock the vault.

        Configures the BW server URL (for self-hosted instances), logs in
        via API key, and unlocks the vault. API key auth does not require
        TOTP -- the key itself is obtained from an authenticated session.
        """
        env = os.environ.copy()
        env["BW_CLIENTID"] = self.client_id
        env["BW_CLIENTSECRET"] = self.client_secret

        # Configure server URL for self-hosted instances
        if self.server_url:
            subprocess.run(
                ["bw", "config", "server", self.server_url],
                capture_output=True, text=True, env=env,
            )

        # Login via API key (tolerates already-logged-in state)
        result = subprocess.run(
            ["bw", "login", "--apikey"],
            capture_output=True, text=True, env=env,
        )
        if result.returncode != 0 and "already" not in result.stderr.lower():
            raise RuntimeError(f"BW login failed: {result.stderr.strip()}")

        # Unlock via password file (more reliable than stdin with native binary)
        with tempfile.NamedTemporaryFile(mode="w", suffix=".pw", delete=False) as pw_file:
            pw_file.write(self.password)
            pw_file_path = pw_file.name
        try:
            self.session = subprocess.run(
                ["bw", "unlock", "--passwordfile", pw_file_path, "--raw"],
                capture_output=True, text=True, env=env,
            ).stdout.strip()
        finally:
            os.unlink(pw_file_path)

        if not self.session:
            raise RuntimeError("BW unlock failed -- no session token returned")

        self.sync()

    def sync(self) -> None:
        """Sync the local vault cache with the server.

        Must be called after login and before any read to guarantee
        the local cache reflects the latest server state. Without this,
        items created by other clients (e.g. the host bw wrapper) will
        not appear in list/search results.
        """
        self._run_bw(["sync"])
        self._last_sync = time.monotonic()

    def _sync_if_stale(self, max_age_s: float = 30.0) -> None:
        """Re-sync if the cache is older than max_age_s seconds.

        Multiple containers share this vault (host wrapper, provisioner
        runs). A read performed on a stale cache sees ghosts: items that
        exist server-side appear missing (or vice versa). Cheap enough
        to run before every read.
        """
        if time.monotonic() - self._last_sync > max_age_s:
            self.sync()

    def _run_bw_with_retry(self, args: list[str], retries: int = 3) -> str:
        """Run a bw command, retrying on transient failures.

        The bw CLI occasionally returns empty output or non-JSON errors
        under load (observed: empty create response, 'Expecting value'
        JSON decode upstream). Retry with backoff before failing.
        """
        last_err = None
        for attempt in range(1, retries + 1):
            try:
                return self._run_bw(args)
            except (RuntimeError, json.JSONDecodeError) as e:
                last_err = e
                msg = str(e)
                # Non-retryable failures: re-raise immediately
                if "not found" in msg.lower() or "already exists" in msg.lower():
                    raise
                if attempt < retries:
                    delay = 2 * attempt
                    time.sleep(delay)
                    self.sync()
        raise RuntimeError(
            f"bw {' '.join(args)} failed after {retries} retries: {last_err}"
        )

    def generate_password(self, length: int = 32) -> str:
        """Generate a strong password."""
        return self._run_bw(["generate", "-ulns", "--length", str(length)])

    def get_totp(self, item_name: str) -> str:
        """Get the current TOTP code for a Bitwarden item.

        Resolves via exact-name match first -- bw's search is fuzzy and
        shared substrings (e.g. "coo" in every tsgstaff-coo-* username)
        make name-based gets ambiguous.
        """
        item_id = self.get_item_id(item_name)
        if not item_id:
            raise RuntimeError(f"Item '{item_name}' not found")
        return self._run_bw(["get", "totp", item_id])

    # -------------------------------------------------------------------
    # Item ID resolution -- the safe way to reference items
    # -------------------------------------------------------------------

    def get_item_id(self, name: str) -> Optional[str]:
        """Resolve an item name to its BW ID.

        Returns the item ID if exactly one match exists, None if no match,
        and raises RuntimeError if multiple items share the name (ambiguous).
        Always syncs if the cache is stale -- containers share this vault
        and a stale cache sees ghosts.
        """
        try:
            self._sync_if_stale()
            output = self._run_bw(["list", "items", "--search", name])
            items = json.loads(output)
            # Filter to exact name matches (bw search is fuzzy)
            matches = [i for i in items if i.get("name") == name]
            if len(matches) == 0:
                return None
            if len(matches) > 1:
                ids = ", ".join(m["id"] for m in matches)
                raise RuntimeError(
                    f"Multiple BW items named '{name}': {ids}. "
                    f"This is a data integrity issue -- resolve manually."
                )
            return matches[0]["id"]
        except json.JSONDecodeError:
            return None

    def get_item(self, name: str) -> Optional[dict]:
        """Get the full item JSON by name. Returns None if not found."""
        item_id = self.get_item_id(name)
        if not item_id:
            return None
        output = self._run_bw(["get", "item", item_id])
        return json.loads(output)

    # -------------------------------------------------------------------
    # Create / Update -- never duplicate, never delete
    # -------------------------------------------------------------------

    def create_item(
        self,
        name: str,
        username: str,
        password: str,
        uris: list[str],
        collection_name: str,
        totp_secret: Optional[str] = None,
        custom_fields: Optional[dict[str, str]] = None,
    ) -> str:
        """Create a login item in a Bitwarden collection.

        Raises RuntimeError if an item with this name already exists.
        Use update_item() to modify an existing item.

        Returns the item ID.
        """
        # SAFEGUARD: refuse to create duplicates
        existing_id = self.get_item_id(name)
        if existing_id:
            raise RuntimeError(
                f"Item '{name}' already exists (id={existing_id}). "
                f"Use update_item() to modify it. "
                f"This safeguard prevents credential duplication."
            )

        item = {
            "type": 1,  # LOGIN
            "name": name,
            "login": {
                "username": username,
                "password": password,
                "uris": [{"uri": u, "match": None} for u in uris],
            },
            "collectionIds": [],
        }

        if totp_secret:
            item["login"]["totp"] = totp_secret

        if custom_fields:
            item["fields"] = [
                {"name": k, "value": v, "type": 0}
                for k, v in custom_fields.items()
            ]

        collection_id = self._get_collection_id(collection_name)
        if collection_id:
            item["collectionIds"] = [collection_id]

        encoded_item = self._encode(item)
        output = self._run_bw_with_retry(["create", "item", encoded_item])
        created = json.loads(output)
        self.sync()
        return created.get("id", "")

    def update_item(
        self,
        name: str,
        password: Optional[str] = None,
        totp_secret: Optional[str] = None,
        uris: Optional[list[str]] = None,
        custom_fields: Optional[dict[str, str]] = None,
    ) -> str:
        """Update an existing item in place by name.

        Only the provided fields are updated; others remain unchanged.
        If the item does not exist, raises RuntimeError.

        Returns the item ID.
        """
        item_id = self.get_item_id(name)
        if not item_id:
            raise RuntimeError(
                f"Cannot update: item '{name}' not found. "
                f"Use create_item() to create it first."
            )

        # Fetch the current item to preserve existing fields
        current = json.loads(self._run_bw(["get", "item", item_id]))

        # Apply updates only to provided fields
        if password is not None:
            current["login"]["password"] = password
        if totp_secret is not None:
            current["login"]["totp"] = totp_secret
        if uris is not None:
            current["login"]["uris"] = [{"uri": u, "match": None} for u in uris]
        if custom_fields is not None:
            current["fields"] = [
                {"name": k, "value": v, "type": 0}
                for k, v in custom_fields.items()
            ]

        encoded_item = self._encode(current)
        output = self._run_bw_with_retry(["edit", "item", item_id, encoded_item])
        updated = json.loads(output)
        self.sync()
        return updated.get("id", item_id)

    # -------------------------------------------------------------------
    # Read helpers
    # -------------------------------------------------------------------

    def item_exists(self, name: str) -> bool:
        """Check if a Bitwarden item with this name already exists."""
        return self.get_item_id(name) is not None

    def get_item_password(self, name: str) -> str:
        """Get the password field from a Bitwarden item."""
        item_id = self.get_item_id(name)
        if not item_id:
            raise RuntimeError(f"Item '{name}' not found")
        return self._run_bw(["get", "password", item_id])

    def get_item_uri(self, name: str) -> str:
        """Get the URI from a Bitwarden item."""
        item = self.get_item(name)
        if not item:
            return ""
        uris = item.get("login", {}).get("uris", [])
        return uris[0]["uri"] if uris else ""

    def list_items(self) -> list[dict]:
        """List all items in the vault."""
        self.sync()
        output = self._run_bw(["list", "items"])
        return json.loads(output)

    # -------------------------------------------------------------------
    # Collections
    # -------------------------------------------------------------------

    def _get_collection_id(self, collection_name: str) -> Optional[str]:
        """Look up a collection ID by name. Returns None if not found."""
        try:
            output = self._run_bw(["list", "collections"])
            collections = json.loads(output)
            for col in collections:
                if col.get("name", "").lower() == collection_name.lower():
                    return col.get("id")
        except (RuntimeError, json.JSONDecodeError):
            pass
        return None

    def create_collection(self, collection_name: str, org_id: str) -> str:
        """Create a collection in an organization."""
        item = {"name": collection_name, "organizationId": org_id}
        encoded_item = self._encode(item)
        output = self._run_bw(["create", "collection", encoded_item])
        created = json.loads(output)
        return created.get("id", "")