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 json
import os
import secrets
import subprocess
import threading
import time
from typing import Callable
from urllib.request import Request, urlopen
from urllib.parse import urlencode


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


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],
        ttl_seconds: int = 900,
        clock: Callable[[], float] = time.monotonic,
    ) -> None:
        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) -> 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,
                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) -> None:
        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)


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:
        request = Request(
            self.base_url + "/auth/simple/login",
            data=json.dumps(
                {"username": "admin", "password": self.admin_password}
            ).encode(),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        with self.opener(request, timeout=10) as response:
            token = str(json.loads(response.read())["token"])
        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()