Setting the file. One moment. 69
Creating Amazon Aurora Db Cluster With Instances
92
Routing Traffic With Route53 And CloudFront
Resilience Program Design
Creating API Gateway Stage
def _exchange_code
— line 169
This file
- Number
- 45.4
- Position
- 4 of 7
- Type
- Python
- Size
- 17 KB
- Lines
- 506
scripts/auth.py
Python·506 lines·17 KB
hashlib
17import html
18import json
19import os
20import secrets
21import signal
22import stat
23import subprocess
24import sys
25import tempfile
26import threading
27import time
28from dataclasses import asdict
29from http.server import BaseHTTPRequestHandler
30from typing import Any, Optional, Tuple
31from urllib.parse import parse_qs, urlencode, urlparse
32
33from launch_config import (
34 AUTH_WAIT_POLL_INTERVAL_SECS,
35 CALLBACK_TIMEOUT_SECS,
36 CLIENT_NAME,
37 DEFAULT_TOKEN_LIFETIME_SECS,
38 ENV_SCOPES,
39 MAX_SESSION_LIFETIME_SECS,
40 SCOPES,
41 SESSION_DIR,
42 SESSION_FILE_NAME,
43 SSO_OIDC_REGION,
44 TOKEN_EXPIRY_BUFFER_SECS,
45 ClientCredentials,
46 StoredSession,
47 resolve_issuer_url,
48)
49
50__version__ = "0.1.0"
51
52
53class SessionExpiredError(Exception):
54 """Raised when the session cannot be refreshed non-interactively."""
55
56 pass
57
58
59def _session_dir() -> str:
60 return os.path.expanduser(SESSION_DIR)
61
62
63def _session_file() -> str:
64 return os.path.join(_session_dir(), SESSION_FILE_NAME)
65
66
67def _create_sso_oidc_client(region: str) -> Any:
68 import boto3
69 from botocore.config import Config as BotoConfig
70
71 config = BotoConfig(user_agent_extra=f"md/awslabs#launch-with-aws#{__version__}")
72 return boto3.client("sso-oidc", region_name=region, config=config)
73
74
75def _generate_pkce() -> Tuple[str, str]:
76 code_verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).rstrip(b"=").decode("ascii")
77 digest = hashlib.sha256(code_verifier.encode("ascii")).digest()
78 code_challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii")
79 return code_verifier, code_challenge
80
81
82def _generate_state() -> str:
83 return secrets.token_hex(16)
84
85
86# ── Session persistence ──────────────────────────────────────────────────
87
88
89def load_session() -> Optional[StoredSession]:
90 try:
91 with open(_session_file()) as f:
92 return StoredSession.from_json(f.read())
93 except (OSError, ValueError, TypeError, KeyError):
94 return None
95
96
97def save_session(session: StoredSession) -> None:
98 directory = _session_dir()
99 os.makedirs(directory, exist_ok=True)
100 os.chmod(directory, stat.S_IRWXU)
101
102 old_umask = os.umask(0o077)
103 try:
104 fd, tmp_path = tempfile.mkstemp(dir=directory, suffix=".tmp")
105 with os.fdopen(fd, "w") as f:
106 f.write(session.to_json(indent=2))
107 os.chmod(tmp_path, stat.S_IRUSR | stat.S_IWUSR)
108 os.replace(tmp_path, _session_file())
109 finally:
110 os.umask(old_umask)
111
112
113def has_valid_session() -> bool:
114 session = load_session()
115 return bool(session and session.token_expires_at > int(time.time()) + TOKEN_EXPIRY_BUFFER_SECS)
116
117
118# ── Client registration ──────────────────────────────────────────────────
119
120
121def _resolve_scopes() -> list[str]:
122 env_val = os.environ.get(ENV_SCOPES)
123 if env_val:
124 return [s.strip() for s in env_val.split(",") if s.strip()]
125 return list(SCOPES)
126
127
128def _register_client() -> ClientCredentials:
129 sso_oidc = _create_sso_oidc_client(SSO_OIDC_REGION)
130 scopes = _resolve_scopes()
131
132 kwargs: dict[str, Any] = {
133 "clientName": CLIENT_NAME,
134 "clientType": "public",
135 "grantTypes": ["authorization_code", "refresh_token"],
136 "scopes": scopes,
137 "issuerUrl": resolve_issuer_url(),
138 "redirectUris": ["http://127.0.0.1"],
139 }
140
141 response = sso_oidc.register_client(**kwargs)
142
143 client_id = response.get("clientId")
144 client_secret = response.get("clientSecret")
145 if not client_id or not client_secret:
146 raise RuntimeError("Client registration response missing credentials")
147
148 # Cap the local session lifetime at MAX_SESSION_LIFETIME_SECS. A
149 # clientSecretExpiresAt of 0 means the IdP set no expiry; otherwise use the
150 # sooner of the IdP value and the cap.
151 expires_at_sec = response.get("clientSecretExpiresAt") or 0
152 now = int(time.time())
153 cap = now + MAX_SESSION_LIFETIME_SECS
154 client_expires_at = cap if expires_at_sec == 0 else min(int(expires_at_sec), cap)
155
156 return ClientCredentials(
157 client_id=client_id,
158 client_secret=client_secret,
159 client_expires_at=client_expires_at,
160 authorize_endpoint=f"https://oidc.{SSO_OIDC_REGION}.amazonaws.com/authorize",
161 token_endpoint=f"https://oidc.{SSO_OIDC_REGION}.amazonaws.com/token",
162 scopes=scopes,
163 )
164
165
166# ── Token exchange ────────────────────────────────────────────────────────
167
168
169def _exchange_code(
170 credentials: ClientCredentials,
171 code: str,
172 code_verifier: str,
173 redirect_uri: str,
174) -> Tuple[str, str, int]:
175 sso_oidc = _create_sso_oidc_client(SSO_OIDC_REGION)
176 response = sso_oidc.create_token(
177 clientId=credentials.client_id,
178 clientSecret=credentials.client_secret,
179 grantType="authorization_code",
180 code=code,
181 codeVerifier=code_verifier,
182 redirectUri=redirect_uri,
183 )
184
185 access_token = response.get("accessToken")
186 refresh_token = response.get("refreshToken")
187 if not access_token or not refresh_token:
188 raise RuntimeError("Token exchange response missing tokens")
189
190 expires_in = response.get("expiresIn") or DEFAULT_TOKEN_LIFETIME_SECS
191 return access_token, refresh_token, int(time.time()) + expires_in
192
193
194def _refresh_token(session: StoredSession) -> Tuple[str, str, int]:
195 sso_oidc = _create_sso_oidc_client(SSO_OIDC_REGION)
196 response = sso_oidc.create_token(
197 clientId=session.client_id,
198 clientSecret=session.client_secret,
199 grantType="refresh_token",
200 refreshToken=session.refresh_token,
201 )
202
203 access_token = response.get("accessToken")
204 if not access_token:
205 raise RuntimeError("Token refresh response missing tokens")
206
207 refresh_token = response.get("refreshToken") or session.refresh_token
208 expires_in = response.get("expiresIn") or DEFAULT_TOKEN_LIFETIME_SECS
209 return access_token, refresh_token, int(time.time()) + expires_in
210
211
212# ── Loopback callback server ─────────────────────────────────────────────
213
214
215class _OAuthState:
216 def __init__(self, expected_state: str) -> None:
217 self.expected_state = expected_state
218 self.auth_code: Optional[str] = None
219 self.error: Optional[str] = None
220 self.code_received = threading.Event()
221
222
223class _CallbackHandler(BaseHTTPRequestHandler):
224 def _respond(self, status: int, message: str) -> None:
225 self.send_response(status)
226 self.send_header("Content-Type", "text/html; charset=utf-8")
227 self.end_headers()
228 body = f'<html><head><meta charset="utf-8"></head><body><h2>{html.escape(message)}</h2></body></html>'
229 self.wfile.write(body.encode("utf-8"))
230
231 def do_GET(self) -> None: # noqa: N802
232 state: _OAuthState = self.server.oauth_state # type: ignore[attr-defined]
233 params = parse_qs(urlparse(self.path).query)
234
235 error = params.get("error", [None])[0]
236 if error:
237 desc = params.get("error_description", [error])[0]
238 self._respond(400, f"Authorization failed: {desc}")
239 state.error = f"Authorization denied: {desc}"
240 state.code_received.set()
241 return
242
243 returned_state = params.get("state", [None])[0]
244 code = params.get("code", [None])[0]
245
246 if returned_state != state.expected_state:
247 if not code:
248 self.send_response(204)
249 self.end_headers()
250 return
251 self._respond(400, "State mismatch")
252 state.error = "OAuth state mismatch - possible CSRF"
253 state.code_received.set()
254 return
255
256 if not code:
257 self.send_response(204)
258 self.end_headers()
259 return
260
261 self._respond(200, "Authenticated - you can close this tab.")
262 state.auth_code = code
263 state.code_received.set()
264
265 def log_message(self, format: str, *args: object) -> None:
266 pass
267
268
269# ── Public API ────────────────────────────────────────────────────────────
270
271
272def get_access_token() -> str:
273 """Return a valid Bearer access token (non-interactive).
274
275 Tries the cached token and silent refresh. If neither works, raises
276 SessionExpiredError so the caller can direct the user to auth-start.
277 """
278 session = load_session()
279 now = int(time.time())
280
281 # 1. Valid cached access token.
282 if session and session.token_expires_at > now + TOKEN_EXPIRY_BUFFER_SECS:
283 return session.access_token
284
285 # 2. Try silent refresh.
286 if session and session.client_expires_at > now:
287 try:
288 access_token, refresh_tok, token_expires_at = _refresh_token(session)
289 updated = StoredSession(
290 **{
291 **asdict(session),
292 "access_token": access_token,
293 "refresh_token": refresh_tok,
294 "token_expires_at": token_expires_at,
295 }
296 )
297 save_session(updated)
298 return updated.access_token
299 except Exception:
300 pass
301
302 raise SessionExpiredError("Session expired or not authenticated. Run auth-start to sign in.")
303
304
305def session_status() -> dict:
306 """Report the local session state without triggering authentication.
307
308 Returns a dict describing whether a session exists, whether its access
309 token is currently usable, and how long until the access token and the
310 overall session registration expire (seconds, clamped at 0).
311 """
312 session = load_session()
313 if not session:
314 return {"authenticated": False}
315
316 now = int(time.time())
317 token_valid = session.token_expires_at > now + TOKEN_EXPIRY_BUFFER_SECS
318 return {
319 "authenticated": token_valid,
320 "tokenExpiresInSecs": max(0, session.token_expires_at - now),
321 "sessionExpiresInSecs": max(0, session.client_expires_at - now),
322 "canRefresh": session.client_expires_at > now,
323 }
324
325
326def sign_out() -> dict:
327 """Sign out: delete the local session so the next call requires re-auth.
328
329 Best-effort revokes the refresh token, then removes the local session file.
330 Revocation failures do not block local deletion.
331 """
332 session = load_session()
333
334 if session and session.refresh_token:
335 try:
336 sso_oidc = _create_sso_oidc_client(SSO_OIDC_REGION)
337 # revoke_token is not exposed on all endpoints; call it only when
338 # available.
339 revoke = getattr(sso_oidc, "revoke_token", None)
340 if callable(revoke):
341 revoke(
342 clientId=session.client_id,
343 clientSecret=session.client_secret,
344 token=session.refresh_token,
345 tokenTypeHint="refresh_token",
346 )
347 except Exception:
348 pass
349
350 removed = False
351 try:
352 os.unlink(_session_file())
353 removed = True
354 except FileNotFoundError:
355 pass
356
357 return {"signedOut": True, "sessionRemoved": removed}
358
359
360# ── Non-blocking auth (auth-start / auth-wait) ──────────────────────────
361
362
363def start_auth() -> dict:
364 """Attempt auth non-blockingly. Returns immediately.
365
366 Returns:
367 - {authenticated: True, reusedCachedSession: True} if cached token valid
368 - {authenticated: True, reusedCachedSession: False} if silent refresh worked
369 - {authenticated: False, signInUrl: ..., pid: N, port: N} if interactive needed
370 """
371 session = load_session()
372 now = int(time.time())
373
374 if session and session.token_expires_at > now + TOKEN_EXPIRY_BUFFER_SECS:
375 return {"authenticated": True, "reusedCachedSession": True}
376
377 if session and session.client_expires_at > now:
378 try:
379 access_token, refresh_tok, token_expires_at = _refresh_token(session)
380 updated = StoredSession(
381 **{
382 **asdict(session),
383 "access_token": access_token,
384 "refresh_token": refresh_tok,
385 "token_expires_at": token_expires_at,
386 }
387 )
388 save_session(updated)
389 return {"authenticated": True, "reusedCachedSession": False}
390 except Exception:
391 pass
392
393 credentials: ClientCredentials
394 if session and session.client_expires_at > now:
395 credentials = session
396 else:
397 credentials = _register_client()
398
399 code_verifier, code_challenge = _generate_pkce()
400 state_value = _generate_state()
401
402 server_script = os.path.join(os.path.dirname(__file__), "auth_callback_server.py")
403 # Configuration values are passed to the child over stdin as a JSON blob.
404 proc = subprocess.Popen(
405 [
406 sys.executable,
407 server_script,
408 "--state",
409 state_value,
410 "--timeout",
411 str(int(CALLBACK_TIMEOUT_SECS)),
412 ],
413 stdout=subprocess.PIPE,
414 stderr=subprocess.PIPE,
415 stdin=subprocess.PIPE,
416 start_new_session=True,
417 )
418
419 # Write the config to the child over stdin, then close it so the child sees EOF.
420 assert proc.stdin is not None
421 proc.stdin.write(
422 json.dumps(
423 {
424 "client_id": credentials.client_id,
425 "client_secret": credentials.client_secret,
426 "client_expires_at": credentials.client_expires_at,
427 "scopes": credentials.scopes,
428 "code_verifier": code_verifier,
429 }
430 ).encode()
431 )
432 proc.stdin.close()
433
434 # Child writes port to stdout once bound.
435 assert proc.stdout is not None
436 assert proc.stderr is not None
437 port_line = proc.stdout.readline().decode().strip()
438 proc.stdout.close()
439 if not port_line.isdigit():
440 proc.terminate()
441 err_output = proc.stderr.read().decode().strip()
442 proc.stderr.close()
443 detail = err_output or port_line or "no output"
444 raise RuntimeError(f"Callback server failed to start: {detail}")
445 proc.stderr.close()
446 port = int(port_line)
447
448 redirect_uri = f"http://127.0.0.1:{port}"
449 authorize_url = f"{credentials.authorize_endpoint}?" + urlencode(
450 {
451 "response_type": "code",
452 "client_id": credentials.client_id,
453 "redirect_uri": redirect_uri,
454 "state": state_value,
455 "code_challenge": code_challenge,
456 "code_challenge_method": "S256",
457 "scope": " ".join(credentials.scopes),
458 }
459 )
460
461 return {
462 "authenticated": False,
463 "signInUrl": authorize_url,
464 "pid": proc.pid,
465 "port": port,
466 }
467
468
469def wait_for_auth(pid: int, timeout: Optional[float] = None) -> dict:
470 """Block until the background callback server completes auth.
471
472 Args:
473 pid: PID of the background callback server (from start_auth).
474 timeout: Max seconds to wait (default: CALLBACK_TIMEOUT_SECS).
475
476 Returns:
477 {authenticated: True} on success.
478
479 Raises:
480 TimeoutError if timeout exceeded.
481 RuntimeError if the server exited without completing auth.
482 """
483 timeout = timeout or CALLBACK_TIMEOUT_SECS
484 deadline = time.time() + timeout
485
486 while time.time() < deadline:
487 if has_valid_session():
488 return {"authenticated": True}
489
490 try:
491 os.kill(pid, 0)
492 except OSError:
493 if has_valid_session():
494 return {"authenticated": True}
495 raise RuntimeError(
496 "Auth callback server exited without completing authentication. "
497 "Run auth-start again."
498 )
499
500 time.sleep(AUTH_WAIT_POLL_INTERVAL_SECS)
501
502 try:
503 os.kill(pid, signal.SIGTERM)
504 except OSError:
505 pass
506 raise TimeoutError(f"Authentication timed out after {int(timeout)} seconds")