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/password_setup.py

@@ -12,13 +12,14 @@ import threading
 import time
 from typing import Callable
 from urllib.request import Request, urlopen
-from urllib.parse import urlencode
+from urllib.parse import urlencode, urlsplit
 
 
 @dataclass(frozen=True)
 class SetupGrant:
     subject: str
     expires_at: float
+    return_to: str = ""
 
 
 class PasswordSetupGrants:
@@ -34,9 +35,18 @@ class PasswordSetupGrants:
         *,
         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
@@ -45,7 +55,7 @@ class PasswordSetupGrants:
         self._subjects: dict[str, str] = {}
         self._lock = threading.Lock()
 
-    def issue(self, subject: str) -> str:
+    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)
@@ -56,6 +66,7 @@ class PasswordSetupGrants:
                 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
@@ -67,7 +78,7 @@ class PasswordSetupGrants:
             grant = self._grants.get(digest)
             return grant is not None and grant.expires_at > self.clock()
 
-    def consume(self, token: str, password: str) -> None:
+    def consume(self, token: str, password: str) -> str:
         if len(password) < 12:
             raise ValueError("password must contain at least 12 characters")
         digest = _digest(token)
@@ -77,6 +88,7 @@ class PasswordSetupGrants:
                 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:

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__":

identity-provisioner/tests/test_company_return_http.py

@@ -0,0 +1,77 @@
+import json
+import pathlib
+import sys
+import threading
+import unittest
+from http.server import ThreadingHTTPServer
+from urllib.error import HTTPError
+from urllib.parse import urlencode, urlsplit, parse_qs
+from urllib.request import Request, urlopen
+
+sys.path.insert(0, str(pathlib.Path(__file__).parents[1]))
+from password_setup import PasswordSetupGrants
+from provisioner import Result
+from server import Handler
+
+
+class CompanyReturnHTTPTests(unittest.TestCase):
+    def setUp(self):
+        self.calls = []
+        class Provisioner:
+            def provision(self, payload):
+                return Result('fixture', 'uid=recipient,ou=people', 'password_setup_required', False)
+        class TestHandler(Handler):
+            service_token = 'fixture-only-service-token'
+            provisioner = Provisioner()
+            password_setups = PasswordSetupGrants(
+                public_url='https://kc.example/setup/password',
+                setter=lambda subject, password: self.calls.append(subject),
+                tenant_returns={'tenant:trial:demo': 'https://app.example/demo/'},
+            )
+        self.handler = TestHandler
+        self.server = ThreadingHTTPServer(('127.0.0.1', 0), TestHandler)
+        self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
+        self.thread.start()
+        self.base = 'http://127.0.0.1:' + str(self.server.server_port)
+
+    def tearDown(self):
+        self.server.shutdown()
+        self.server.server_close()
+        self.thread.join()
+
+    def issue(self):
+        payload = {'user_id': 'fixture', 'tenant': 'tenant:trial:demo',
+                   'primary_email': 'fixture@example.test', 'idempotency_key': 'fixture-long-idempotency',
+                   'correlation_id': 'fixture', 'return_to': 'https://evil.example/'}
+        request = Request(self.base + '/v1/identities/provision', data=json.dumps(payload).encode(),
+                          headers={'Authorization': 'Bearer fixture-only-service-token'})
+        with urlopen(request) as response:
+            url = json.load(response)['password_setup_url']
+        return parse_qs(urlsplit(url).query)['token'][0]
+
+    def submit(self, token):
+        request = Request(self.base + '/setup/password', data=urlencode({
+            'token': token, 'password': 'fixture-long-password', 'confirmation': 'fixture-long-password',
+            'return_to': 'https://evil.example/',
+        }).encode(), headers={'Content-Type': 'application/x-www-form-urlencoded'})
+        try:
+            response = urlopen(request)
+        except HTTPError as error:
+            response = error
+        with response:
+            return response.status, response.read().decode(), response.headers
+
+    def test_authenticated_issue_and_completion_ignore_browser_return(self):
+        token = self.issue()
+        status, body, headers = self.submit(token)
+        self.assertEqual(200, status)
+        self.assertIn('https://app.example/demo/', body)
+        self.assertNotIn('evil.example', body)
+        self.assertNotIn(token, body)
+        self.assertNotIn('recipient', body)
+        self.assertEqual('no-referrer', headers['Referrer-Policy'])
+        self.assertEqual(['recipient'], self.calls)
+        status, body, headers = self.submit(token)
+        self.assertEqual(400, status)
+        self.assertNotIn('https://app.example/demo/', body)
+        self.assertEqual(['recipient'], self.calls)

identity-provisioner/tests/test_password_setup.py

@@ -47,3 +47,38 @@ class PasswordSetupGrantTests(unittest.TestCase):
         with self.assertRaisesRegex(ValueError, "12 characters"):
             self.grants.consume(token, "too-short")
         self.assertTrue(self.grants.valid(token))
+
+
+class CompanyReturnTests(unittest.TestCase):
+    def test_return_is_bound_to_grant_and_cannot_be_changed_by_browser(self):
+        mapping = {"tenant:trial:demo-company": "https://vergabe.example/demo-company/"}
+        grants = PasswordSetupGrants(public_url="https://kc.example/setup/password",
+                                     setter=lambda *args: None, tenant_returns=mapping)
+        url = grants.issue("recipient", "tenant:trial:demo-company")
+        token = url.partition("token=")[2]
+        mapping["tenant:trial:demo-company"] = "https://attacker.example/"
+        self.assertNotIn("recipient", url)
+        self.assertNotIn("return", url)
+        self.assertEqual("https://vergabe.example/demo-company/", grants.consume(token, "test-password-long"))
+        with self.assertRaises(ValueError):
+            grants.consume(token, "test-password-long")
+
+    def test_unknown_tenant_has_no_return_and_bad_targets_fail(self):
+        for target in ("http://example.test/", "https://example.test/?next=evil",
+                       "https://example.test/#fragment", "https://user:password@example.test/"):
+            with self.assertRaises(ValueError):
+                PasswordSetupGrants(public_url="https://kc.example/setup/password",
+                                    setter=lambda *args: None, tenant_returns={"tenant:trial:demo": target})
+        grants = PasswordSetupGrants(public_url="https://kc.example/setup/password", setter=lambda *args: None)
+        token = grants.issue("recipient", "tenant:unknown").partition("token=")[2]
+        self.assertEqual("", grants.consume(token, "test-password-long"))
+
+    def test_expiry_and_failure_never_release_a_return(self):
+        clock = [0]
+        grants = PasswordSetupGrants(public_url="https://kc.example/setup/password",
+                                    setter=lambda *args: None, clock=lambda: clock[0], ttl_seconds=1,
+                                    tenant_returns={"tenant:trial:demo": "https://app.example/demo/"})
+        token = grants.issue("recipient", "tenant:trial:demo").partition("token=")[2]
+        clock[0] = 2
+        with self.assertRaises(ValueError):
+            grants.consume(token, "test-password-long")