ariadne/ariadne/services/vaultwarden_sync.py

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timezone
import time
from typing import Any

from ..settings import settings
from ..utils.logging import get_logger
from .keycloak_admin import keycloak_admin
from .mailu import mailu
from .vaultwarden import vaultwarden


VAULTWARDEN_EMAIL_ATTR = "vaultwarden_email"
VAULTWARDEN_STATUS_ATTR = "vaultwarden_status"
VAULTWARDEN_SYNCED_AT_ATTR = "vaultwarden_synced_at"
VAULTWARDEN_MASTER_ATTR = "vaultwarden_master_password_set_at"

logger = get_logger(__name__)


@dataclass(frozen=True)
class VaultwardenSyncSummary:
    processed: int
    created_or_present: int
    skipped: int
    failures: int
    detail: str = ""


@dataclass
class VaultwardenSyncCounters:
    processed: int = 0
    created_or_present: int = 0
    skipped: int = 0
    failures: int = 0

    def summary(self, detail: str = "") -> VaultwardenSyncSummary:
        return VaultwardenSyncSummary(
            processed=self.processed,
            created_or_present=self.created_or_present,
            skipped=self.skipped,
            failures=self.failures,
            detail=detail,
        )

    def status(self) -> str:
        return "ok" if self.failures == 0 else "error"


@dataclass
class VaultwardenInviteState:
    username: str
    status: str
    synced_at: str
    synced_ts: float | None
    full_user: dict[str, Any]
    counters: VaultwardenSyncCounters


def _extract_attr(attrs: Any, key: str) -> str:
    if not isinstance(attrs, dict):
        return ""
    raw = attrs.get(key)
    if isinstance(raw, list):
        for item in raw:
            if isinstance(item, str) and item.strip():
                return item.strip()
        return ""
    if isinstance(raw, str) and raw.strip():
        return raw.strip()
    return ""


def _parse_synced_at(value: str) -> float | None:
    value = (value or "").strip()
    if not value:
        return None
    for fmt in ("%Y-%m-%dT%H:%M:%SZ", "%Y-%m-%dT%H:%M:%S%z"):
        try:
            parsed = datetime.strptime(value, fmt)
            if parsed.tzinfo is None:
                parsed = parsed.replace(tzinfo=timezone.utc)
            return parsed.timestamp()
        except ValueError:
            continue
    return None


def _vaultwarden_email_for_user(user: dict[str, Any]) -> str:
    username = (user.get("username") if isinstance(user.get("username"), str) else "") or ""
    username = username.strip()
    if not username:
        return ""

    attrs = user.get("attributes")
    vaultwarden_email = _extract_attr(attrs, VAULTWARDEN_EMAIL_ATTR)
    if vaultwarden_email:
        return vaultwarden_email

    mailu_email = _extract_attr(attrs, "mailu_email")
    if mailu_email:
        return mailu_email

    email = (user.get("email") if isinstance(user.get("email"), str) else "") or ""
    email = email.strip()
    if email and email.lower().endswith(f"@{settings.mailu_domain.lower()}"):
        return email

    return f"{username}@{settings.mailu_domain}"


def _set_user_attribute_if_missing(username: str, user: dict[str, Any], key: str, value: str) -> None:
    value = (value or "").strip()
    if not value:
        return
    existing = _extract_attr(user.get("attributes"), key)
    if existing:
        return
    keycloak_admin.set_user_attribute(username, key, value)


def _set_user_attribute(username: str, key: str, value: str) -> None:
    value = (value or "").strip()
    if not value:
        return
    keycloak_admin.set_user_attribute(username, key, value)


def _normalize_user(user: dict[str, Any]) -> tuple[str, dict[str, Any]] | None:
    username = (user.get("username") if isinstance(user.get("username"), str) else "") or ""
    username = username.strip()
    if not username:
        return None
    if user.get("enabled") is False:
        return None
    if user.get("serviceAccountClientId") or username.startswith("service-account-"):
        return None

    user_id = (user.get("id") if isinstance(user.get("id"), str) else "") or ""
    full_user = user
    if user_id:
        try:
            full_user = keycloak_admin.get_user(user_id)
        except Exception:
            full_user = user
    return username, full_user


def _current_sync_state(full_user: dict[str, Any]) -> tuple[str, str, float | None]:
    current_status = _extract_attr(full_user.get("attributes"), VAULTWARDEN_STATUS_ATTR)
    current_synced_at = _extract_attr(full_user.get("attributes"), VAULTWARDEN_SYNCED_AT_ATTR)
    current_synced_ts = _parse_synced_at(current_synced_at)
    return current_status, current_synced_at, current_synced_ts


def _cooldown_active(status: str, synced_ts: float | None) -> bool:
    if status not in {"rate_limited", "error"} or not synced_ts:
        return False
    return time.time() - synced_ts < settings.vaultwarden_retry_cooldown_sec


def _has_pending_failures(users: list[dict[str, Any]]) -> bool:
    for user in users:
        username = (user.get("username") if isinstance(user.get("username"), str) else "") or ""
        username = username.strip()
        if not username or user.get("enabled") is False:
            continue
        if user.get("serviceAccountClientId") or username.startswith("service-account-"):
            continue
        attrs = user.get("attributes") if isinstance(user.get("attributes"), dict) else {}
        status = _extract_attr(attrs, VAULTWARDEN_STATUS_ATTR)
        synced_at = _extract_attr(attrs, VAULTWARDEN_SYNCED_AT_ATTR)
        master_set_at = _extract_attr(attrs, VAULTWARDEN_MASTER_ATTR)
        synced_ts = _parse_synced_at(synced_at)
        if not status:
            return True
        if status in {"invited", "already_present"} and not synced_at:
            return True
        if status == "already_present" and not master_set_at:
            return True
        if status == "invited" and _should_refresh_invite(synced_ts):
            return True
        if status in {"error", "rate_limited"} and not _cooldown_active(status, synced_ts):
            return True
    return False


def _set_sync_status(username: str, status: str) -> None:
    try:
        _set_user_attribute(username, VAULTWARDEN_STATUS_ATTR, status)
        _set_user_attribute(
            username,
            VAULTWARDEN_SYNCED_AT_ATTR,
            time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        )
    except Exception:
        return


def _set_master_password_set(username: str, full_user: dict[str, Any]) -> None:
    if _extract_attr(full_user.get("attributes"), VAULTWARDEN_MASTER_ATTR):
        return
    try:
        _set_user_attribute(
            username,
            VAULTWARDEN_MASTER_ATTR,
            time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        )
    except Exception:
        return


def _ensure_email_attrs(username: str, full_user: dict[str, Any], email: str) -> None:
    try:
        _set_user_attribute_if_missing(username, full_user, "mailu_email", email)
        _set_user_attribute_if_missing(username, full_user, VAULTWARDEN_EMAIL_ATTR, email)
    except Exception:
        return


def _should_refresh_invite(synced_ts: float | None) -> bool:
    if settings.vaultwarden_invite_refresh_sec <= 0:
        return False
    if synced_ts is None:
        return True
    return (time.time() - synced_ts) >= settings.vaultwarden_invite_refresh_sec


def _handle_existing_invite(state: VaultwardenInviteState) -> bool:
    if state.status not in {"invited", "already_present"}:
        return False
    if state.status == "already_present":
        if not state.synced_at:
            _set_sync_status(state.username, state.status)
        _set_master_password_set(state.username, state.full_user)
        state.counters.skipped += 1
        return True
    if not state.synced_at:
        _set_sync_status(state.username, state.status)
        state.counters.skipped += 1
        return True
    if not _should_refresh_invite(state.synced_ts):
        if not state.synced_at:
            _set_sync_status(state.username, state.status)
        state.counters.skipped += 1
        return True
    return False


def _sync_user(
    user: dict[str, Any],
    counters: VaultwardenSyncCounters,
) -> tuple[str | None, bool]:
    status: str | None = None
    ok = False
    normalized = _normalize_user(user)
    if not normalized:
        counters.skipped += 1
    else:
        username, full_user = normalized
        current_status, current_synced_at, current_synced_ts = _current_sync_state(full_user)
        if _cooldown_active(current_status, current_synced_ts):
            counters.skipped += 1
        else:
            email = _vaultwarden_email_for_user(full_user)
            if not email:
                counters.skipped += 1
            elif not mailu.mailbox_exists(email):
                counters.skipped += 1
            else:
                _ensure_email_attrs(username, full_user, email)
                state = VaultwardenInviteState(
                    username=username,
                    status=current_status,
                    synced_at=current_synced_at,
                    synced_ts=current_synced_ts,
                    full_user=full_user,
                    counters=counters,
                )
                if _handle_existing_invite(state):
                    status = None
                else:
                    counters.processed += 1
                    result = vaultwarden.invite_user(email)
                    status = result.status
                    if result.ok:
                        counters.created_or_present += 1
                        ok = True
                    else:
                        counters.failures += 1
                    _set_sync_status(username, result.status)
                    if result.status == "already_present":
                        _set_master_password_set(username, full_user)
    return status, ok


def run_vaultwarden_sync() -> VaultwardenSyncSummary:
    consecutive_failures = 0
    counters = VaultwardenSyncCounters()

    if not keycloak_admin.ready():
        counters.failures = 1
        summary = counters.summary(detail="keycloak admin not configured")
        logger.info(
            "vaultwarden sync skipped",
            extra={"event": "vaultwarden_sync", "status": "error", "detail": summary.detail},
        )
        return summary

    users = list(keycloak_admin.iter_users(page_size=200, brief=False))
    if not _has_pending_failures(users):
        summary = counters.summary(detail="no pending failures")
        logger.info(
            "vaultwarden sync skipped",
            extra={"event": "vaultwarden_sync", "status": "skip", "detail": summary.detail},
        )
        return summary

    for user in users:
        status, ok = _sync_user(user, counters)
        if status is None:
            continue
        if ok:
            consecutive_failures = 0
            continue
        if status in {"rate_limited", "error"}:
            consecutive_failures += 1
        if consecutive_failures >= settings.vaultwarden_failure_bailout:
            break

    summary = counters.summary()
    logger.info(
        "vaultwarden sync finished",
        extra={
            "event": "vaultwarden_sync",
            "status": counters.status(),
            "processed": counters.processed,
            "created_or_present": counters.created_or_present,
            "skipped": counters.skipped,
            "failures": counters.failures,
        },
    )
    return summary