ariadne/ariadne/services/vaultwarden.py

from __future__ import annotations

import threading
import time
from dataclasses import dataclass

import httpx

from ..k8s.client import get_secret_value, get_json
from ..settings import settings


HTTP_OK = 200
HTTP_CREATED = 201
HTTP_NO_CONTENT = 204
HTTP_BAD_REQUEST = 400
HTTP_CONFLICT = 409
HTTP_TOO_MANY_REQUESTS = 429


@dataclass(frozen=True)
class VaultwardenInvite:
    ok: bool
    status: str
    detail: str = ""


class VaultwardenService:
    def __init__(self) -> None:
        self._admin_lock = threading.Lock()
        self._admin_client: httpx.Client | None = None
        self._admin_session_expires_at: float = 0.0
        self._admin_session_base_url: str = ""
        self._rate_limited_until: float = 0.0

    @staticmethod
    def _normalize_email(email: str) -> str | None:
        email = (email or "").strip()
        if not email or "@" not in email:
            return None
        return email

    def _rate_limited(self) -> VaultwardenInvite:
        return VaultwardenInvite(ok=False, status="rate_limited", detail="vaultwarden rate limited")

    def _candidate_urls(self) -> list[str]:
        base_url = f"http://{settings.vaultwarden_service_host}"
        urls = [base_url]
        try:
            pod_ip = self._find_pod_ip(settings.vaultwarden_namespace, settings.vaultwarden_pod_label)
            urls.append(f"http://{pod_ip}:{settings.vaultwarden_pod_port}")
        except Exception:
            pass
        return [url for url in urls if url]

    @staticmethod
    def _already_present(resp: httpx.Response) -> bool:
        try:
            body = resp.text or ""
        except Exception:
            body = ""
        lowered = body.lower()
        return any(
            marker in lowered
            for marker in (
                "already invited",
                "already exists",
                "already registered",
                "user already exists",
            )
        )

    def _invite_via(self, base_url: str, email: str) -> VaultwardenInvite | None:
        if not base_url:
            return None
        try:
            session = self._admin_session(base_url)
            resp = session.post("/admin/invite", json={"email": email})
            if resp.status_code == HTTP_TOO_MANY_REQUESTS:
                self._rate_limited_until = time.time() + float(settings.vaultwarden_admin_rate_limit_backoff_sec)
                result = self._rate_limited()
            elif resp.status_code in {HTTP_OK, HTTP_CREATED, HTTP_NO_CONTENT}:
                result = VaultwardenInvite(ok=True, status="invited", detail="invite created")
            elif resp.status_code in {HTTP_BAD_REQUEST, HTTP_CONFLICT} and self._already_present(resp):
                result = VaultwardenInvite(ok=True, status="already_present", detail="user already present")
            else:
                result = VaultwardenInvite(ok=False, status="error", detail=f"status {resp.status_code}")
        except Exception as exc:
            message = str(exc)
            if "rate limited" in message.lower():
                result = self._rate_limited()
            else:
                result = VaultwardenInvite(ok=False, status="error", detail=message)
        return result

    def invite_user(self, email: str) -> VaultwardenInvite:
        email = self._normalize_email(email)
        if not email:
            return VaultwardenInvite(ok=False, status="invalid_email", detail="email invalid")
        if self._rate_limited_until and time.time() < self._rate_limited_until:
            return self._rate_limited()

        last_error = ""
        for candidate in self._candidate_urls():
            result = self._invite_via(candidate, email)
            if not result:
                continue
            if result.ok:
                return result
            if result.status == "rate_limited":
                return result
            last_error = result.detail or last_error

        return VaultwardenInvite(ok=False, status="error", detail=last_error or "failed to invite")

    def _admin_session(self, base_url: str) -> httpx.Client:
        now = time.time()
        with self._admin_lock:
            if self._rate_limited_until and now < self._rate_limited_until:
                raise RuntimeError("vaultwarden rate limited")
            if self._admin_client and now < self._admin_session_expires_at and self._admin_session_base_url == base_url:
                return self._admin_client

            if self._admin_client:
                try:
                    self._admin_client.close()
                except Exception:
                    pass
                self._admin_client = None

            token = get_secret_value(
                settings.vaultwarden_namespace,
                settings.vaultwarden_admin_secret_name,
                settings.vaultwarden_admin_secret_key,
            )

            client = httpx.Client(
                base_url=base_url,
                timeout=10.0,
                follow_redirects=True,
                headers={"User-Agent": "ariadne/1"},
            )
            resp = client.post("/admin", data={"token": token})
            if resp.status_code == HTTP_TOO_MANY_REQUESTS:
                self._rate_limited_until = now + float(settings.vaultwarden_admin_rate_limit_backoff_sec)
                raise RuntimeError("vaultwarden rate limited")
            resp.raise_for_status()

            self._admin_client = client
            self._admin_session_base_url = base_url
            self._admin_session_expires_at = now + float(settings.vaultwarden_admin_session_ttl_sec)
            return client

    @staticmethod
    def _find_pod_ip(namespace: str, label_selector: str) -> str:
        data = get_json(f"/api/v1/namespaces/{namespace}/pods?labelSelector={label_selector}")
        items = data.get("items") or []
        if not isinstance(items, list) or not items:
            raise RuntimeError("no vaultwarden pods found")

        def _pod_ready(pod: dict) -> bool:
            status = pod.get("status") if isinstance(pod.get("status"), dict) else {}
            if status.get("phase") != "Running":
                return False
            ip = status.get("podIP")
            if not isinstance(ip, str) or not ip:
                return False
            conditions = status.get("conditions") if isinstance(status.get("conditions"), list) else []
            for cond in conditions:
                if not isinstance(cond, dict):
                    continue
                if cond.get("type") == "Ready":
                    return cond.get("status") == "True"
            return True

        ready = [p for p in items if isinstance(p, dict) and _pod_ready(p)]
        candidates = ready or [p for p in items if isinstance(p, dict)]
        status = candidates[0].get("status") or {}
        ip = status.get("podIP") if isinstance(status, dict) else None
        if not isinstance(ip, str) or not ip:
            raise RuntimeError("vaultwarden pod has no IP")
        return ip


vaultwarden = VaultwardenService()