Coverage for src / quber / playground / session.py: 73%
173 statements
« prev ^ index » next coverage.py v7.14.0, created at 2026-09-23 22:14 -0400
1"""The hosted playground's sign-in through WorkOS, and the session cookie it
2leaves behind.
4Two programs have to agree on this: the app's login gate, and the wake
5function that stands in for the app while no task is running. Whichever of
6the two the load balancer routes a request to has to be able to complete a
7sign-in, so both serve the same three routes from the functions here, and
8the app accepts a cookie the function set without a second sign-in. The wake
9function's package is built from this file and its own, so everything here
10is standard library only, the call to WorkOS included.
12WorkOS proves who a person is and which organization they belong to. The
13playground decides who may enter: only members of the one organization named
14in settings. The WorkOS access token is read once, when the sign-in
15completes, and thrown away; what the playground keeps is its own cookie.
17A cookie value and a sign-in ``state`` share one format: a JSON record,
18base64url encoded, a dot, and an HMAC-SHA256 over the encoded record under
19the session secret. The record's ``kind`` keeps the two apart, so neither can
20stand in for the other. A cookie's expiry is fixed when it is issued; using
21the playground does not extend it. When it lapses while the WorkOS session is
22still alive, the person goes to WorkOS and straight back without typing
23anything.
25``/login`` answers with a page rather than a bare redirect. A link scraper
26never signs in, so the page a pasted link unfurls from is this one, and it
27carries the preview tags; a browser moves on to WorkOS at once.
28"""
30from __future__ import annotations
32import base64
33import hashlib
34import hmac
35import json
36import time
37import urllib.error
38import urllib.request
39from dataclasses import asdict, dataclass
40from html import escape
41from typing import Any, Optional
42from urllib.parse import urlencode
44COOKIE = "quber_session"
46# What a link to the playground unfurls into when it is pasted into a chat or
47# a social post. A link scraper never signs in, so the page it reads is the
48# sign-in page, and that page is where these tags live. The card image has to
49# be reachable without a session in both states of the service, so the app
50# serves it from its static files and the wake function serves the same
51# bytes from its own package.
52SITE_URL = "https://playground.qubera.ai"
53LINK_CARD_PATH = "/static/link-card.png"
54LINK_CARD_WIDTH = 1200
55LINK_CARD_HEIGHT = 630
56LINK_TITLE = "Qubera — extraction playground"
57LINK_DESCRIPTION = (
58 "Upload a financial filing and ask a question. Every figure in the answer "
59 "is cited to the page and the table cell it came from."
60)
61SESSION_HOURS = 2
62SESSION_SECONDS = SESSION_HOURS * 3600
63# How long a sign-in may take between leaving for WorkOS and coming back.
64STATE_SECONDS = 15 * 60
65WORKOS_API = "https://api.workos.com"
66WORKOS_TIMEOUT_SECONDS = 8
69@dataclass(frozen=True)
70class WorkOS:
71 """What both programs need to run a sign-in: the WorkOS environment's API
72 key and client ID, the one organization whose members may enter, and the
73 secret that signs the cookie and the ``state``."""
75 api_key: str
76 client_id: str
77 organization_id: str
78 session_secret: str
81@dataclass(frozen=True)
82class Identity:
83 """Who a session belongs to, as WorkOS reported it at sign-in. ``sid`` is
84 the WorkOS session id, which sign-out needs to end that session too."""
86 user_id: str
87 email: str
88 organization_id: str
89 role: str
90 sid: str
93@dataclass(frozen=True)
94class SignIn:
95 """How a return from WorkOS ends. Accepted: a redirect to the page the
96 person asked for, with a cookie. Refused or failed: a page, no cookie."""
98 status: int
99 location: Optional[str] = None
100 token: Optional[str] = None
101 page: Optional[str] = None
102 identity: Optional[Identity] = None
105class SignInError(Exception):
106 """WorkOS did not confirm the sign-in."""
109def signature(secret: str, body: str) -> str:
110 return hmac.new(secret.encode(), body.encode(), hashlib.sha256).hexdigest()
113def b64decode(text: str) -> bytes:
114 return base64.urlsafe_b64decode(text + "=" * (-len(text) % 4))
117def encode(secret: str, record: dict[str, Any]) -> str:
118 body = base64.urlsafe_b64encode(json.dumps(record, separators=(",", ":")).encode()).decode().rstrip("=")
119 return f"{body}.{signature(secret, body)}"
122def decode(
123 secret: str, token: Optional[str], kind: str, now: Optional[float] = None
124) -> Optional[dict[str, Any]]:
125 """The record inside ``token`` when its signature holds, it is of ``kind``,
126 and it has not expired. ``now=None`` checks against the clock; pass
127 ``float("-inf")`` to ignore the expiry."""
128 if not token or token.count(".") != 1:
129 return None
130 body, given = token.split(".")
131 if not hmac.compare_digest(signature(secret, body), given):
132 return None
133 try:
134 record = json.loads(b64decode(body))
135 except ValueError:
136 return None
137 if not isinstance(record, dict) or record.get("kind") != kind:
138 return None
139 expires = record.get("expires")
140 if not isinstance(expires, int) or expires <= (now if now is not None else time.time()):
141 return None
142 return record
145def issue(secret: str, identity: Identity, now: Optional[float] = None) -> str:
146 expires = int(now if now is not None else time.time()) + SESSION_SECONDS
147 return encode(secret, {"kind": "session", **asdict(identity), "expires": expires})
150def identity_of(record: Optional[dict[str, Any]]) -> Optional[Identity]:
151 if record is None:
152 return None
153 try:
154 return Identity(
155 user_id=str(record["user_id"]),
156 email=str(record["email"]),
157 organization_id=str(record["organization_id"]),
158 role=str(record["role"]),
159 sid=str(record["sid"]),
160 )
161 except KeyError:
162 return None
165def verify(
166 secret: str, organization_id: str, token: Optional[str], now: Optional[float] = None
167) -> Optional[Identity]:
168 """The identity in a live session cookie for ``organization_id``, or None.
169 A cookie issued under another organization is refused like an expired one."""
170 identity = identity_of(decode(secret, token, "session", now))
171 if identity is None or not hmac.compare_digest(identity.organization_id, organization_id):
172 return None
173 return identity
176def signed_identity(secret: str, token: Optional[str]) -> Optional[Identity]:
177 """The identity in a cookie whose signature holds, expired or not. Sign-out
178 reads the WorkOS session id this way, so a person whose cookie has just
179 lapsed can still end their WorkOS session."""
180 return identity_of(decode(secret, token, "session", float("-inf")))
183def state_for(secret: str, next_path: Optional[str], now: Optional[float] = None) -> str:
184 expires = int(now if now is not None else time.time()) + STATE_SECONDS
185 return encode(secret, {"kind": "state", "next": safe_next(next_path), "expires": expires})
188def next_from_state(secret: str, state: Optional[str], now: Optional[float] = None) -> Optional[str]:
189 """The return page a ``state`` carries, or None when it is forged or stale."""
190 record = decode(secret, state, "state", now)
191 if record is None:
192 return None
193 return safe_next(str(record.get("next", "")))
196def safe_next(raw: Optional[str]) -> str:
197 """The page to land on after sign-in: a path on this site, never elsewhere."""
198 if raw and raw.startswith("/") and not raw.startswith("//") and "\\" not in raw:
199 return raw
200 return "/"
203def base_url(scheme: str, host: str) -> str:
204 """The site a request came in on. The callback and the sign-out return
205 both hang off it, so the hosted playground and a developer machine each
206 get their own; WorkOS accepts only the ones registered in its dashboard."""
207 return f"{scheme}://{host}"
210def authorize_url(config: WorkOS, base: str, next_path: Optional[str], now: Optional[float] = None) -> str:
211 query = urlencode(
212 {
213 "client_id": config.client_id,
214 "redirect_uri": base + "/callback",
215 "response_type": "code",
216 "provider": "authkit",
217 "state": state_for(config.session_secret, next_path, now),
218 }
219 )
220 return f"{WORKOS_API}/user_management/authorize?{query}"
223def logout_url(sid: str, base: str) -> str:
224 """Where to send a browser to end its WorkOS session. WorkOS then returns it
225 to ``base``, which has to be a registered sign-out URI."""
226 query = urlencode({"session_id": sid, "return_to": base + "/"})
227 return f"{WORKOS_API}/user_management/sessions/logout?{query}"
230def authenticate(config: WorkOS, code: str) -> dict[str, Any]:
231 """Trade the code WorkOS sent back for the person's identity."""
232 body = json.dumps(
233 {
234 "client_id": config.client_id,
235 "client_secret": config.api_key,
236 "grant_type": "authorization_code",
237 "code": code,
238 }
239 ).encode()
240 request = urllib.request.Request(
241 f"{WORKOS_API}/user_management/authenticate",
242 data=body,
243 headers={"Content-Type": "application/json", "Accept": "application/json"},
244 method="POST",
245 )
246 try:
247 with urllib.request.urlopen(request, timeout=WORKOS_TIMEOUT_SECONDS) as response:
248 answer = json.loads(response.read())
249 except urllib.error.HTTPError as exc:
250 detail = exc.read().decode(errors="replace")[:300]
251 raise SignInError(f"WorkOS answered {exc.code}: {detail}") from exc
252 except (urllib.error.URLError, TimeoutError, ValueError) as exc:
253 raise SignInError(f"WorkOS could not be reached: {exc}") from exc
254 if not isinstance(answer, dict):
255 raise SignInError("WorkOS answered with something other than an object")
256 return answer
259def claims(access_token: str) -> dict[str, Any]:
260 """The claims inside a WorkOS access token, read without checking its
261 signature. That is safe here only because the token arrived straight from
262 WorkOS over HTTPS in the same call."""
263 parts = access_token.split(".")
264 if len(parts) != 3:
265 return {}
266 try:
267 found = json.loads(b64decode(parts[1]))
268 except ValueError:
269 return {}
270 return found if isinstance(found, dict) else {}
273def complete_sign_in(
274 config: WorkOS,
275 base: str,
276 code: Optional[str],
277 state: Optional[str],
278 mark_src: str,
279 now: Optional[float] = None,
280) -> SignIn:
281 """Finish a sign-in WorkOS sent back to ``/callback``.
283 The organization check happens here: an account outside the configured
284 organization is refused, and gets no cookie. WorkOS has already signed
285 that account in by then, so the refusal page offers to end the WorkOS
286 session; otherwise the next sign-in would come straight back as the same
287 account until WorkOS times it out.
288 """
289 target = next_from_state(config.session_secret, state, now)
290 if target is None or not code:
291 return SignIn(400, page=failed_page(mark_src))
292 try:
293 answer = authenticate(config, code)
294 except SignInError:
295 return SignIn(400, page=failed_page(mark_src))
296 raw_user = answer.get("user")
297 user: dict[str, Any] = raw_user if isinstance(raw_user, dict) else {}
298 token_claims = claims(str(answer.get("access_token", "")))
299 email = str(user.get("email", ""))
300 sid = str(token_claims.get("sid", ""))
301 if answer.get("organization_id") != config.organization_id:
302 switch = logout_url(sid, base) if sid else "/login"
303 return SignIn(403, page=refusal_page(email, switch, mark_src))
304 identity = Identity(
305 user_id=str(user.get("id", "")),
306 email=email,
307 organization_id=config.organization_id,
308 role=str(token_claims.get("role", "")),
309 sid=sid,
310 )
311 return SignIn(303, location=target, token=issue(config.session_secret, identity, now), identity=identity)
314def cookie_header(token: str, *, secure: bool) -> str:
315 """The Set-Cookie value for a session, as the app and the wake function
316 both send it: HttpOnly, SameSite=Lax, and Secure on HTTPS, which is every
317 hosted request."""
318 return cookie_value(token, SESSION_SECONDS, secure=secure)
321def clear_cookie_header(*, secure: bool) -> str:
322 return cookie_value("", 0, secure=secure)
325def cookie_value(token: str, max_age: int, *, secure: bool) -> str:
326 parts = [f"{COOKIE}={token}", f"Max-Age={max_age}", "Path=/", "HttpOnly", "SameSite=Lax"]
327 if secure:
328 parts.append("Secure")
329 return "; ".join(parts)
332_PAGE = """<!DOCTYPE html>
333<html lang="en">
334<head>
335<meta charset="utf-8" />
336<meta name="viewport" content="width=device-width, initial-scale=1" />
337<title>{title}</title>
338{head}
339<link rel="preconnect" href="https://fonts.googleapis.com" />
340<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin />
341<link rel="stylesheet" href="https://fonts.googleapis.com/css2?family=IBM+Plex+Mono:wght@400;500&family=IBM+Plex+Sans:wght@400;600&family=IBM+Plex+Serif:wght@400;600&display=swap" />
342<style>
343 /* Values follow the playground's shared token set. */
344 * {{ box-sizing: border-box; }}
345 body {{ margin: 0; min-height: 100vh; display: grid; place-items: center; padding: 24px 16px;
346 background: #f5f8fb; color: #15222e;
347 font-family: "IBM Plex Sans", system-ui, -apple-system, "Segoe UI", sans-serif;
348 -webkit-font-smoothing: antialiased; }}
349 main {{ width: min(420px, 100%); padding: 36px; background: #ffffff; border: 1px solid #e2e8ef;
350 border-radius: 14px; box-shadow: 0 4px 14px rgba(12, 32, 50, .08), 0 18px 50px rgba(12, 32, 50, .10);
351 display: flex; flex-direction: column; gap: 18px; }}
352 .brand {{ display: flex; align-items: center; gap: 12px; }}
353 .brand img {{ width: 32px; height: 32px; display: block; }}
354 .brand span {{ font-family: "IBM Plex Serif", Georgia, serif; font-weight: 600; font-size: 22px;
355 letter-spacing: -0.01em; }}
356 .copy {{ display: flex; flex-direction: column; gap: 10px; }}
357 .eyebrow {{ margin: 0; font-family: "IBM Plex Mono", ui-monospace, monospace; font-size: 11px;
358 letter-spacing: 0.15em; text-transform: uppercase; color: #74828f; }}
359 h1 {{ margin: 0; font-family: "IBM Plex Serif", Georgia, serif; font-weight: 400; font-size: 26px;
360 line-height: 1.15; letter-spacing: -0.015em; }}
361 p.message {{ margin: 0; font-size: 15px; line-height: 1.6; color: #475563; overflow-wrap: anywhere; }}
362 a.action {{ align-self: flex-start; min-height: 44px; padding: 0 18px; display: inline-flex; align-items: center;
363 border-radius: 999px; background: #1E4E78; color: #ffffff; font-size: 14px; font-weight: 600;
364 text-decoration: none; }}
365 a.action:hover {{ background: #163b5c; }}
366 a.action:focus-visible {{ outline: none; box-shadow: 0 0 0 3px rgba(191, 139, 58, .25); }}
367</style>
368</head>
369<body>
370<main>
371 <div class="brand"><img src="{mark}" alt="" /><span>Qubera</span></div>
372 <div class="copy">
373 <p class="eyebrow">{eyebrow}</p>
374 <h1>{heading}</h1>
375 <p class="message">{message}</p>
376 </div>
377 <a class="action" href="{href}">{action}</a>
378</main>
379</body>
380</html>
381"""
384def page(
385 *,
386 title: str,
387 eyebrow: str,
388 heading: str,
389 message: str,
390 action: str,
391 href: str,
392 mark_src: str,
393 head: str = "",
394) -> str:
395 """One card on the paper background, with the Qubera mark and wordmark,
396 in the design system's type and colors. Self-contained apart from the
397 fonts and ``mark_src``, because the wake function serves it while nothing
398 under the app's /static/ answers; each program passes a mark it can serve."""
399 return _PAGE.format(
400 title=escape(title),
401 head=head,
402 mark=escape(mark_src, quote=True),
403 eyebrow=escape(eyebrow),
404 heading=escape(heading),
405 message=escape(message),
406 href=escape(href, quote=True),
407 action=escape(action),
408 )
411def login_page(config: WorkOS, base: str, next_path: Optional[str], mark_src: str) -> str:
412 """The page ``/login`` answers with: preview tags for link scrapers, and an
413 immediate move to WorkOS for a browser."""
414 target = authorize_url(config, base, next_path)
415 refresh = f'<meta http-equiv="refresh" content="0; url={escape(target, quote=True)}" />'
416 return page(
417 title="Qubera — sign in",
418 eyebrow="Sign in",
419 heading="Taking you to sign in",
420 message="The playground is for invited members. Your browser should move on by itself.",
421 action="Continue to sign in",
422 href=target,
423 mark_src=mark_src,
424 head=link_preview_tags() + "\n" + refresh,
425 )
428def refusal_page(email: str, switch_href: str, mark_src: str) -> str:
429 who = email or "This account"
430 return page(
431 title="Qubera — no access",
432 eyebrow="Access",
433 heading="This account can’t open the playground",
434 message=f"{who} is not a member of the Qubera organization. Only members can sign in.",
435 action="Sign in with a different account",
436 href=switch_href,
437 mark_src=mark_src,
438 )
441def failed_page(mark_src: str) -> str:
442 return page(
443 title="Qubera — sign-in did not finish",
444 eyebrow="Sign in",
445 heading="The sign-in didn’t finish",
446 message="It could not be confirmed, or it took longer than 15 minutes. Start it again.",
447 action="Sign in again",
448 href="/login",
449 mark_src=mark_src,
450 )
453def link_preview_tags() -> str:
454 """The Open Graph and Twitter Card tags, one per line, for a page's head."""
455 image = SITE_URL + LINK_CARD_PATH
456 title = escape(LINK_TITLE, quote=True)
457 description = escape(LINK_DESCRIPTION, quote=True)
458 tags = [
459 ("name", "description", description),
460 ("property", "og:type", "website"),
461 ("property", "og:site_name", "Qubera"),
462 ("property", "og:url", SITE_URL + "/"),
463 ("property", "og:title", title),
464 ("property", "og:description", description),
465 ("property", "og:image", image),
466 ("property", "og:image:secure_url", image),
467 ("property", "og:image:type", "image/png"),
468 ("property", "og:image:width", str(LINK_CARD_WIDTH)),
469 ("property", "og:image:height", str(LINK_CARD_HEIGHT)),
470 ("property", "og:image:alt", title),
471 ("name", "twitter:card", "summary_large_image"),
472 ("name", "twitter:title", title),
473 ("name", "twitter:description", description),
474 ("name", "twitter:image", image),
475 ("name", "twitter:image:alt", title),
476 ]
477 return "\n".join(f'<meta {attr}="{key}" content="{value}" />' for attr, key, value in tags)