net-kingdom/identity-provisioner/server.py

from dataclasses import asdict
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from html import escape
import json
import os
from urllib.parse import parse_qs, urlsplit
import secrets

from provisioner import LLDAPProvisioner, _directory_username, dispatch
from password_setup import LLDAPPasswordSetter, PasswordSetupGrants


class Handler(BaseHTTPRequestHandler):
    provisioner: LLDAPProvisioner
    password_setups: PasswordSetupGrants
    service_token: str

    def do_GET(self):
        path = urlsplit(self.path)
        if path.path == "/healthz":
            return self._send(200, {"status": "ok"})
        if path.path == "/setup/password":
            token = parse_qs(path.query).get("token", [""])[0]
            if not self.password_setups.valid(token):
                return self._html(400, _expired_page())
            return self._html(200, _setup_page(token))
        self._send(404, {"error": "not_found"})

    def do_POST(self):
        path = urlsplit(self.path).path
        if path == "/setup/password":
            return self._set_password()
        supplied = self.headers.get("Authorization", "").removeprefix("Bearer ")
        if not secrets.compare_digest(supplied, self.service_token):
            return self._send(403, {"error": "access_denied"})
        try:
            length = min(int(self.headers.get("Content-Length", "0")), 65536)
            payload = json.loads(self.rfile.read(length))
            result = dispatch(self.provisioner, path, payload)
        except KeyError:
            return self._send(404, {"error": "not_found"})
        except (ValueError, json.JSONDecodeError) as exc:
            return self._send(400, {"error": "invalid_request", "message": str(exc)})
        except RuntimeError:
            return self._send(503, {"error": "dependency_unavailable"})
        response = asdict(result)
        if path == "/v1/identities/provision" and result.status == "password_setup_required":
            response["password_setup_url"] = self.password_setups.issue(
                _directory_username(result.external_subject)
            )
        self._send(200, response)

    def _set_password(self):
        if self.headers.get("Content-Type", "").partition(";")[0] != "application/x-www-form-urlencoded":
            return self._html(400, _expired_page())
        length = min(int(self.headers.get("Content-Length", "0")), 8192)
        body = parse_qs(self.rfile.read(length).decode("utf-8", "replace"))
        token = body.get("token", [""])[0]
        password = body.get("password", [""])[0]
        confirmation = body.get("confirmation", [""])[0]
        if password != confirmation:
            return self._html(400, _setup_page(token, "Passwords do not match."))
        try:
            self.password_setups.consume(token, password)
        except ValueError as exc:
            return self._html(400, _setup_page(token, str(exc)))
        except RuntimeError:
            return self._html(503, _failed_page())
        return self._html(200, _complete_page())

    def _send(self, status: int, payload: dict):
        body = json.dumps(payload, separators=(",", ":")).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Cache-Control", "no-store")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _html(self, status: int, body: str):
        payload = body.encode()
        self.send_response(status)
        self.send_header("Content-Type", "text/html; charset=utf-8")
        self.send_header("Cache-Control", "no-store")
        self.send_header("Content-Security-Policy", "default-src 'none'; style-src 'unsafe-inline'; form-action 'self'; frame-ancestors 'none'; base-uri 'none'")
        self.send_header("Referrer-Policy", "no-referrer")
        self.send_header("X-Content-Type-Options", "nosniff")
        self.send_header("X-Frame-Options", "DENY")
        self.send_header("Content-Length", str(len(payload)))
        self.end_headers()
        self.wfile.write(payload)

    def log_message(self, format, *args):
        return


def main():
    Handler.provisioner = LLDAPProvisioner(
        base_url=os.environ["LLDAP_URL"],
        admin_password=os.environ["LLDAP_ADMIN_PASSWORD"],
    )
    Handler.service_token = os.environ["PROVISIONER_SERVICE_TOKEN"].strip()
    if not Handler.service_token:
        raise ValueError("PROVISIONER_SERVICE_TOKEN must not be empty")
    Handler.password_setups = PasswordSetupGrants(
        public_url=os.environ["PASSWORD_SETUP_PUBLIC_URL"],
        setter=LLDAPPasswordSetter(
            base_url=os.environ["LLDAP_URL"],
            admin_password=os.environ["LLDAP_ADMIN_PASSWORD"],
        ),
        ttl_seconds=int(os.environ.get("PASSWORD_SETUP_TTL_SECONDS", "900")),
    )
    ThreadingHTTPServer(("0.0.0.0", 8080), Handler).serve_forever()


def _page(title: str, content: str) -> str:
    return f"""<!doctype html><html lang="en"><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>{escape(title)}</title><style>
body{{font:16px system-ui;max-width:36rem;margin:4rem auto;padding:0 1rem}}
label,input,button{{display:block;width:100%;box-sizing:border-box;margin:.75rem 0}}
input,button{{padding:.7rem}}.error{{color:#a00}}
</style><main><h1>{escape(title)}</h1>{content}</main></html>"""


def _setup_page(token: str, error: str = "") -> str:
    message = f'<p class="error">{escape(error)}</p>' if error else ""
    return _page("Set up your password", f"""{message}
<p>This single-use link expires shortly. Choose a password of at least 12 characters.</p>
<form method="post" action="/setup/password" autocomplete="off">
<input type="hidden" name="token" value="{escape(token)}">
<label>New password<input type="password" name="password" minlength="12" required autocomplete="new-password"></label>
<label>Confirm password<input type="password" name="confirmation" minlength="12" required autocomplete="new-password"></label>
<button type="submit">Set password</button></form>""")


def _expired_page() -> str:
    return _page("Link unavailable", "<p>This password setup link is invalid or expired. Request a new link from your tenant administrator.</p>")


def _failed_page() -> str:
    return _page("Setup unavailable", "<p>Password setup could not be completed. Request a new link and try again.</p>")


def _complete_page() -> str:
    return _page("Password set", "<p>Your password is ready. Return to the application to sign in and enroll MFA.</p>")


if __name__ == "__main__":
    main()