Bind password setup grants to approved company welcome pages

Assistant: codex
Assistant-Model: gpt-6-astra
Assistant-Session: 01a092fe-13b1-7f12-ac74-7d258af4d79c

identity-provisioner/server.py

@@ -46,7 +46,7 @@ class Handler(BaseHTTPRequestHandler):
         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)
+                _directory_username(result.external_subject), tenant=str(payload.get("tenant", ""))
             )
         self._send(200, response)
 
@@ -61,12 +61,12 @@ class Handler(BaseHTTPRequestHandler):
         if password != confirmation:
             return self._html(400, _setup_page(token, "Passwords do not match."))
         try:
-            self.password_setups.consume(token, password)
+            return_to = 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())
+        return self._html(200, _complete_page(return_to))
 
     def _send(self, status: int, payload: dict):
         body = json.dumps(payload, separators=(",", ":")).encode()
@@ -108,6 +108,7 @@ def main():
             base_url=os.environ["LLDAP_URL"],
             admin_password=os.environ["LLDAP_ADMIN_PASSWORD"],
         ),
+        tenant_returns=json.loads(os.environ.get("PASSWORD_SETUP_TENANT_RETURNS", "{}")),
         ttl_seconds=int(os.environ.get("PASSWORD_SETUP_TTL_SECONDS", "900")),
     )
     ThreadingHTTPServer(("0.0.0.0", 8080), Handler).serve_forever()
@@ -142,8 +143,13 @@ 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>")
+def _complete_page(return_to: str = "") -> str:
+    content = "<p>Your password is ready. Sign in with your own account to continue.</p>"
+    if return_to:
+        content += f'<p><a rel="noreferrer" href="{escape(return_to, quote=True)}">Continue to your company</a></p>'
+    else:
+        content += "<p>Return to the application to sign in and enroll MFA.</p>"
+    return _page("Password set", content)
 
 
 if __name__ == "__main__":