net-kingdom/identity-provisioner/password_setup.py

"""Short-lived, single-use password setup for managed LLDAP identities."""

from __future__ import annotations

from dataclasses import dataclass
import hashlib
import os
import secrets
import subprocess
import threading
import time
from typing import Callable
from urllib.parse import urlencode, urlsplit
from urllib.request import urlopen

from provisioner import directory_login


@dataclass(frozen=True)
class SetupGrant:
    subject: str
    expires_at: float
    return_to: str = ""


class PasswordSetupGrants:
    """In-memory, fail-closed grants.

    A restart invalidates every outstanding link. Issuing a new link for one
    subject revokes that subject's previous link, and consumption removes the
    grant before the password helper runs so a link can never be replayed.
    """

    def __init__(
        self,
        *,
        public_url: str,
        setter: Callable[[str, str], None],
        tenant_returns: dict[str, str] | None = None,
        ttl_seconds: int = 900,
        clock: Callable[[], float] = time.monotonic,
    ) -> None:
        self.tenant_returns = dict(tenant_returns or {})
        for tenant, target in self.tenant_returns.items():
            parts = urlsplit(target)
            if (not tenant.startswith("tenant:") or parts.scheme != "https"
                    or not parts.hostname or parts.username or parts.password
                    or parts.query or parts.fragment or not parts.path.startswith("/")
                    or parts.hostname in {"localhost", "127.0.0.1"}):
                raise ValueError("tenant welcome targets require exact HTTPS URLs")
        self.public_url = public_url.rstrip("/")
        self.setter = setter
        self.ttl_seconds = ttl_seconds
        self.clock = clock
        self._grants: dict[str, SetupGrant] = {}
        self._subjects: dict[str, str] = {}
        self._lock = threading.Lock()

    def issue(self, subject: str, tenant: str = "") -> str:
        if not subject or len(subject) > 255:
            raise ValueError("valid external subject is required")
        token = secrets.token_urlsafe(32)
        digest = _digest(token)
        with self._lock:
            previous = self._subjects.pop(subject, None)
            if previous is not None:
                self._grants.pop(previous, None)
            self._grants[digest] = SetupGrant(
                subject=subject,
                return_to=self.tenant_returns.get(tenant, ""),
                expires_at=self.clock() + self.ttl_seconds,
            )
            self._subjects[subject] = digest
        return self.public_url + "?" + urlencode({"token": token})

    def valid(self, token: str) -> bool:
        digest = _digest(token)
        with self._lock:
            grant = self._grants.get(digest)
            return grant is not None and grant.expires_at > self.clock()

    def consume(self, token: str, password: str) -> str:
        if len(password) < 12:
            raise ValueError("password must contain at least 12 characters")
        digest = _digest(token)
        with self._lock:
            grant = self._grants.pop(digest, None)
            if grant is None or grant.expires_at <= self.clock():
                raise ValueError("password setup link is invalid or expired")
            self._subjects.pop(grant.subject, None)
        self.setter(grant.subject, password)
        return grant.return_to


class LLDAPPasswordSetter:
    """Invoke LLDAP's official OPAQUE registration helper.

    The user password is passed only through ``LLDAP_USER_PASSWORD``. The
    helper has no environment option for its administrative credential, so we
    first exchange the long-lived admin password for a short-lived LLDAP JWT
    and pass only that token inside the pod's isolated process namespace.
    """

    def __init__(
        self,
        *,
        base_url: str,
        admin_password: str,
        helper: str = "/app/lldap_set_password",
        opener: Callable = urlopen,
        runner: Callable = subprocess.run,
    ) -> None:
        self.base_url = base_url.rstrip("/")
        self.admin_password = admin_password
        self.helper = helper
        self.opener = opener
        self.runner = runner

    def __call__(self, subject: str, password: str) -> None:
        token = directory_login(
            base_url=self.base_url,
            admin_password=self.admin_password,
            opener=self.opener,
            timeout=10,
        )
        env = dict(os.environ)
        env["LLDAP_USER_PASSWORD"] = password
        result = self.runner(
            [
                self.helper,
                "--base-url",
                self.base_url,
                "--token",
                token,
                "--username",
                subject,
            ],
            env=env,
            capture_output=True,
            text=True,
            timeout=20,
            check=False,
        )
        if result.returncode != 0:
            raise RuntimeError("LLDAP password setup failed")


def _digest(token: str) -> str:
    return hashlib.sha256(token.encode()).hexdigest()