| 1 | """SSO business logic — Azure Entra ID, Google & Cloudflare Access.""" |
| 2 | |
| 3 | import datetime |
| 4 | import hashlib |
| 5 | import hmac |
| 6 | import os |
| 7 | import re |
| 8 | import secrets |
| 9 | from typing import Optional |
| 10 | from urllib.parse import urlencode |
| 11 | |
| 12 | import httpx |
| 13 | import jwt |
| 14 | from jwt.algorithms import RSAAlgorithm |
| 15 | from loguru import logger |
| 16 | from sqlalchemy.ext.asyncio import AsyncSession |
| 17 | from sqlmodel import select |
| 18 | |
| 19 | from app.auth.models.sso import SSOAllowedEmail |
| 20 | from app.auth.models.sso import SSOConfig |
| 21 | from app.auth.models.users import User |
| 22 | from app.auth.utils import AuthHandler |
| 23 | from app.db.db_session import async_engine |
| 24 | |
| 25 | # ── HMAC-signed OAuth2 state (stateless, multi-worker safe) ────────────────── |
| 26 | # State format: "<nonce>:<unix_ts>:<hmac_hex>" |
| 27 | # Self-validating — no in-memory or DB store needed. Works across multiple |
| 28 | # uvicorn workers and survives process restarts. |
| 29 | _STATE_TTL_SECONDS = 600 |
| 30 | # Use a dedicated SSO_STATE_SECRET if provided; fall back to JWT_SECRET so |
| 31 | # existing deployments that haven't set this variable continue to work. |
| 32 | _STATE_SECRET = os.environ.get("SSO_STATE_SECRET") or AuthHandler().secret |
| 33 | |
| 34 | # Allowed pattern for Cloudflare team domains — must end in cloudflareaccess.com |
| 35 | _CF_DOMAIN_RE = re.compile(r"^[a-zA-Z0-9-]+\.cloudflareaccess\.com$") |
| 36 | |
| 37 | |
| 38 | def _generate_state() -> str: |
| 39 | """Generate a self-validating HMAC-signed OAuth2 state token.""" |
| 40 | nonce = secrets.token_urlsafe(16) |
| 41 | ts = str(int(datetime.datetime.utcnow().timestamp())) |
| 42 | msg = f"{nonce}:{ts}".encode() |
| 43 | mac = hmac.new(_STATE_SECRET.encode(), msg, hashlib.sha256).hexdigest() |
| 44 | return f"{nonce}:{ts}:{mac}" |
| 45 | |
| 46 | |
| 47 | def _validate_state(state: str) -> bool: |
| 48 | """Validate a self-signed state token. Returns True if valid and not expired.""" |
| 49 | try: |
| 50 | nonce, ts_str, mac = state.rsplit(":", 2) |
| 51 | msg = f"{nonce}:{ts_str}".encode() |
| 52 | expected = hmac.new(_STATE_SECRET.encode(), msg, hashlib.sha256).hexdigest() |
| 53 | if not hmac.compare_digest(mac, expected): |
| 54 | return False |
| 55 | age = datetime.datetime.utcnow().timestamp() - float(ts_str) |
| 56 | return 0 <= age <= _STATE_TTL_SECONDS |
| 57 | except Exception: |
| 58 | return False |
| 59 | |
| 60 | |
| 61 | # ── SSO Config CRUD ────────────────────────────────────────────────────────── |
| 62 | |
| 63 | |
| 64 | async def get_sso_config() -> Optional[SSOConfig]: |
| 65 | async with AsyncSession(async_engine) as session: |
| 66 | result = await session.execute(select(SSOConfig).where(SSOConfig.id == 1)) |
| 67 | return result.scalars().first() |
| 68 | |
| 69 | |
| 70 | async def upsert_sso_config(data: dict) -> SSOConfig: |
| 71 | async with AsyncSession(async_engine) as session: |
| 72 | result = await session.execute(select(SSOConfig).where(SSOConfig.id == 1)) |
| 73 | cfg = result.scalars().first() |
| 74 | if cfg is None: |
| 75 | cfg = SSOConfig(id=1) |
| 76 | session.add(cfg) |
| 77 | for key, value in data.items(): |
| 78 | if not hasattr(cfg, key): |
| 79 | continue |
| 80 | # Booleans: always set (False is a valid value) |
| 81 | # Strings/None: set if key present in payload — allows explicit clearing |
| 82 | # Skip provider secrets if None (don't overwrite existing secrets) |
| 83 | if key in ("azure_client_secret", "google_client_secret") and value is None: |
| 84 | continue |
| 85 | setattr(cfg, key, value) |
| 86 | cfg.updated_at = datetime.datetime.utcnow() |
| 87 | session.add(cfg) |
| 88 | await session.commit() |
| 89 | await session.refresh(cfg) |
| 90 | return cfg |
| 91 | |
| 92 | |
| 93 | # ── Allowed Emails CRUD ────────────────────────────────────────────────────── |
| 94 | |
| 95 | |
| 96 | async def list_allowed_emails() -> list[SSOAllowedEmail]: |
| 97 | async with AsyncSession(async_engine) as session: |
| 98 | result = await session.execute(select(SSOAllowedEmail).order_by(SSOAllowedEmail.id)) |
| 99 | return list(result.scalars().all()) |
| 100 | |
| 101 | |
| 102 | async def add_allowed_email(email: str, role_id: int = 2) -> SSOAllowedEmail: |
| 103 | email = email.lower().strip() |
| 104 | async with AsyncSession(async_engine) as session: |
| 105 | # Check duplicate |
| 106 | result = await session.execute(select(SSOAllowedEmail).where(SSOAllowedEmail.email == email)) |
| 107 | existing = result.scalars().first() |
| 108 | if existing: |
| 109 | raise ValueError(f"Email {email} is already in the allowlist") |
| 110 | entry = SSOAllowedEmail(email=email, role_id=role_id) |
| 111 | session.add(entry) |
| 112 | await session.commit() |
| 113 | await session.refresh(entry) |
| 114 | return entry |
| 115 | |
| 116 | |
| 117 | async def delete_allowed_email(email_id: int) -> bool: |
| 118 | async with AsyncSession(async_engine) as session: |
| 119 | result = await session.execute(select(SSOAllowedEmail).where(SSOAllowedEmail.id == email_id)) |
| 120 | entry = result.scalars().first() |
| 121 | if entry is None: |
| 122 | return False |
| 123 | await session.delete(entry) |
| 124 | await session.commit() |
| 125 | return True |
| 126 | |
| 127 | |
| 128 | async def find_allowed_email(email: str) -> Optional[SSOAllowedEmail]: |
| 129 | async with AsyncSession(async_engine) as session: |
| 130 | result = await session.execute(select(SSOAllowedEmail).where(SSOAllowedEmail.email == email)) |
| 131 | return result.scalars().first() |
| 132 | |
| 133 | |
| 134 | async def find_user_by_email(email: str) -> Optional[User]: |
| 135 | """Look up an existing user by email address.""" |
| 136 | async with AsyncSession(async_engine) as session: |
| 137 | result = await session.execute(select(User).where(User.email == email)) |
| 138 | return result.scalars().first() |
| 139 | |
| 140 | |
| 141 | # ── Auto‑provision SSO user ────────────────────────────────────────────────── |
| 142 | |
| 143 | |
| 144 | async def get_or_create_sso_user(email: str, role_id: int = 2) -> User: |
| 145 | """Find existing user by email or create a new SSO‑managed user.""" |
| 146 | import bcrypt |
| 147 | |
| 148 | async with AsyncSession(async_engine) as session: |
| 149 | result = await session.execute(select(User).where(User.email == email)) |
| 150 | user = result.scalars().first() |
| 151 | if user: |
| 152 | return user |
| 153 | |
| 154 | # Create a new user with a random unusable password. |
| 155 | # token_urlsafe(48) → 64 chars, safely under bcrypt's 72-byte limit. |
| 156 | random_pw = secrets.token_urlsafe(48) |
| 157 | hashed = bcrypt.hashpw(random_pw.encode("utf-8"), bcrypt.gensalt()).decode("utf-8") |
| 158 | |
| 159 | username = email.split("@")[0] |
| 160 | # Ensure unique username |
| 161 | base = username |
| 162 | counter = 1 |
| 163 | while True: |
| 164 | result = await session.execute(select(User).where(User.username == username)) |
| 165 | if result.scalars().first() is None: |
| 166 | break |
| 167 | username = f"{base}{counter}" |
| 168 | counter += 1 |
| 169 | |
| 170 | user = User( |
| 171 | username=username, |
| 172 | password=hashed, |
| 173 | email=email, |
| 174 | role_id=role_id, |
| 175 | ) |
| 176 | session.add(user) |
| 177 | await session.commit() |
| 178 | await session.refresh(user) |
| 179 | logger.info(f"SSO: Auto‑provisioned user '{username}' ({email}) with role_id={role_id}") |
| 180 | return user |
| 181 | |
| 182 | |
| 183 | # ── Azure Entra ID helpers ─────────────────────────────────────────────────── |
| 184 | |
| 185 | _azure_keys_cache: dict[str, tuple[datetime.datetime, list]] = {} |
| 186 | |
| 187 | |
| 188 | def build_azure_auth_url(cfg: SSOConfig) -> str: |
| 189 | """Build the Azure OAuth2 authorization URL.""" |
| 190 | state = _generate_state() |
| 191 | params = { |
| 192 | "client_id": cfg.azure_client_id, |
| 193 | "response_type": "code", |
| 194 | "redirect_uri": cfg.azure_redirect_uri, |
| 195 | "response_mode": "query", |
| 196 | "scope": "openid profile email", |
| 197 | "state": state, |
| 198 | "nonce": secrets.token_urlsafe(16), |
| 199 | "prompt": "select_account", |
| 200 | } |
| 201 | base = f"https://login.microsoftonline.com/{cfg.azure_tenant_id}/oauth2/v2.0/authorize" |
| 202 | return f"{base}?{urlencode(params)}" |
| 203 | |
| 204 | |
| 205 | async def exchange_azure_code(code: str, state: str, cfg: SSOConfig) -> dict: |
| 206 | """Exchange an authorization code for tokens and return the ID token claims.""" |
| 207 | if not _validate_state(state): |
| 208 | raise ValueError("Invalid or expired OAuth2 state parameter") |
| 209 | |
| 210 | token_url = f"https://login.microsoftonline.com/{cfg.azure_tenant_id}/oauth2/v2.0/token" |
| 211 | data = { |
| 212 | "client_id": cfg.azure_client_id, |
| 213 | "client_secret": cfg.azure_client_secret, |
| 214 | "code": code, |
| 215 | "redirect_uri": cfg.azure_redirect_uri, |
| 216 | "grant_type": "authorization_code", |
| 217 | "scope": "openid profile email", |
| 218 | } |
| 219 | |
| 220 | async with httpx.AsyncClient(timeout=15) as client: |
| 221 | resp = await client.post(token_url, data=data) |
| 222 | resp.raise_for_status() |
| 223 | token_data = resp.json() |
| 224 | |
| 225 | id_token = token_data.get("id_token") |
| 226 | if not id_token: |
| 227 | raise ValueError("No id_token in Azure response") |
| 228 | |
| 229 | # Fetch JWKS for signature validation (cached, with retry on key miss) |
| 230 | jwks_url = f"https://login.microsoftonline.com/{cfg.azure_tenant_id}/discovery/v2.0/keys" |
| 231 | header = jwt.get_unverified_header(id_token) |
| 232 | kid = header.get("kid") |
| 233 | |
| 234 | async def _get_azure_keys() -> list: |
| 235 | now = datetime.datetime.utcnow() |
| 236 | cache_key = cfg.azure_tenant_id |
| 237 | if cache_key in _azure_keys_cache: |
| 238 | cached_at, keys = _azure_keys_cache[cache_key] |
| 239 | if (now - cached_at).total_seconds() < 3600: |
| 240 | return keys |
| 241 | async with httpx.AsyncClient(timeout=10) as client: |
| 242 | resp = await client.get(jwks_url) |
| 243 | resp.raise_for_status() |
| 244 | jwks = resp.json() |
| 245 | keys = [{"kid": k.get("kid"), "key": RSAAlgorithm.from_jwk(k)} for k in jwks.get("keys", [])] |
| 246 | _azure_keys_cache[cache_key] = (now, keys) |
| 247 | return keys |
| 248 | |
| 249 | public_key = None |
| 250 | for attempt in range(2): |
| 251 | for k in await _get_azure_keys(): |
| 252 | if k["kid"] == kid: |
| 253 | public_key = k["key"] |
| 254 | break |
| 255 | if public_key: |
| 256 | break |
| 257 | _azure_keys_cache.pop(cfg.azure_tenant_id, None) |
| 258 | |
| 259 | if public_key is None: |
| 260 | raise ValueError("Unable to find matching signing key in Azure JWKS") |
| 261 | |
| 262 | claims = jwt.decode( |
| 263 | id_token, |
| 264 | key=public_key, |
| 265 | algorithms=["RS256"], |
| 266 | audience=cfg.azure_client_id, |
| 267 | issuer=f"https://login.microsoftonline.com/{cfg.azure_tenant_id}/v2.0", |
| 268 | options={"verify_exp": True, "verify_aud": True, "verify_iss": True}, |
| 269 | ) |
| 270 | |
| 271 | return claims |
| 272 | |
| 273 | |
| 274 | # ── Google OAuth2 / OIDC helpers ──────────────────────────────────────────── |
| 275 | |
| 276 | _GOOGLE_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth" |
| 277 | _GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token" |
| 278 | _GOOGLE_JWKS_URL = "https://www.googleapis.com/oauth2/v3/certs" |
| 279 | _GOOGLE_ISSUER = "https://accounts.google.com" |
| 280 | |
| 281 | # Cache for Google public keys (same pattern as Cloudflare) |
| 282 | _google_keys_cache: dict[str, tuple[datetime.datetime, list]] = {} |
| 283 | |
| 284 | |
| 285 | def build_google_auth_url(cfg: SSOConfig) -> str: |
| 286 | """Build the Google OAuth2 authorization URL.""" |
| 287 | state = _generate_state() |
| 288 | params = { |
| 289 | "client_id": cfg.google_client_id, |
| 290 | "response_type": "code", |
| 291 | "redirect_uri": cfg.google_redirect_uri, |
| 292 | "scope": "openid email", |
| 293 | "state": state, |
| 294 | "access_type": "online", |
| 295 | "prompt": "select_account", |
| 296 | } |
| 297 | return f"{_GOOGLE_AUTH_URL}?{urlencode(params)}" |
| 298 | |
| 299 | |
| 300 | async def _get_google_public_keys() -> list: |
| 301 | """Fetch and cache Google's public JWKS keys (TTL 1 hour).""" |
| 302 | now = datetime.datetime.utcnow() |
| 303 | if "google" in _google_keys_cache: |
| 304 | cached_at, keys = _google_keys_cache["google"] |
| 305 | if (now - cached_at).total_seconds() < 3600: |
| 306 | return keys |
| 307 | |
| 308 | async with httpx.AsyncClient(timeout=10) as client: |
| 309 | resp = await client.get(_GOOGLE_JWKS_URL) |
| 310 | resp.raise_for_status() |
| 311 | data = resp.json() |
| 312 | |
| 313 | keys = [] |
| 314 | for jwk in data.get("keys", []): |
| 315 | public_key = RSAAlgorithm.from_jwk(jwk) |
| 316 | keys.append({"kid": jwk.get("kid"), "key": public_key}) |
| 317 | |
| 318 | _google_keys_cache["google"] = (now, keys) |
| 319 | return keys |
| 320 | |
| 321 | |
| 322 | async def exchange_google_code(code: str, state: str, cfg: SSOConfig) -> dict: |
| 323 | """Exchange a Google authorization code for tokens and return ID token claims.""" |
| 324 | if not _validate_state(state): |
| 325 | raise ValueError("Invalid or expired OAuth2 state parameter") |
| 326 | |
| 327 | data = { |
| 328 | "client_id": cfg.google_client_id, |
| 329 | "client_secret": cfg.google_client_secret, |
| 330 | "code": code, |
| 331 | "redirect_uri": cfg.google_redirect_uri, |
| 332 | "grant_type": "authorization_code", |
| 333 | } |
| 334 | |
| 335 | async with httpx.AsyncClient(timeout=15) as client: |
| 336 | resp = await client.post(_GOOGLE_TOKEN_URL, data=data) |
| 337 | resp.raise_for_status() |
| 338 | token_data = resp.json() |
| 339 | |
| 340 | id_token = token_data.get("id_token") |
| 341 | if not id_token: |
| 342 | raise ValueError("No id_token in Google response") |
| 343 | |
| 344 | # Fetch JWKS and find matching key; retry once on miss (handles key rotation) |
| 345 | header = jwt.get_unverified_header(id_token) |
| 346 | kid = header.get("kid") |
| 347 | |
| 348 | for attempt in range(2): |
| 349 | keys = await _get_google_public_keys() |
| 350 | for k in keys: |
| 351 | if k["kid"] == kid: |
| 352 | claims = jwt.decode( |
| 353 | id_token, |
| 354 | key=k["key"], |
| 355 | algorithms=["RS256"], |
| 356 | audience=cfg.google_client_id, |
| 357 | issuer=_GOOGLE_ISSUER, |
| 358 | options={"verify_exp": True, "verify_aud": True, "verify_iss": True}, |
| 359 | ) |
| 360 | return claims |
| 361 | # Key not found — bust the cache and retry once |
| 362 | _google_keys_cache.pop("google", None) |
| 363 | |
| 364 | raise ValueError("Unable to find matching signing key in Google JWKS") |
| 365 | |
| 366 | |
| 367 | # ── Cloudflare Access helpers ──────────────────────────────────────────────── |
| 368 | |
| 369 | # Cache for Cloudflare public keys |
| 370 | _cf_keys_cache: dict[str, tuple[datetime.datetime, list]] = {} |
| 371 | |
| 372 | |
| 373 | async def _get_cf_public_keys(team_domain: str) -> list: |
| 374 | """Fetch and cache Cloudflare Access public keys.""" |
| 375 | cache_key = team_domain |
| 376 | now = datetime.datetime.utcnow() |
| 377 | if cache_key in _cf_keys_cache: |
| 378 | cached_at, keys = _cf_keys_cache[cache_key] |
| 379 | if (now - cached_at).total_seconds() < 3600: # cache for 1 hour |
| 380 | return keys |
| 381 | |
| 382 | certs_url = f"https://{team_domain}/cdn-cgi/access/certs" |
| 383 | async with httpx.AsyncClient(timeout=10) as client: |
| 384 | resp = await client.get(certs_url) |
| 385 | resp.raise_for_status() |
| 386 | data = resp.json() |
| 387 | |
| 388 | keys = [] |
| 389 | for jwk in data.get("keys", []): |
| 390 | public_key = RSAAlgorithm.from_jwk(jwk) |
| 391 | keys.append({"kid": jwk.get("kid"), "key": public_key}) |
| 392 | |
| 393 | _cf_keys_cache[cache_key] = (now, keys) |
| 394 | return keys |
| 395 | |
| 396 | |
| 397 | async def validate_cf_jwt(token: str, cfg: SSOConfig) -> dict: |
| 398 | """Validate a Cloudflare Access JWT assertion and return its claims.""" |
| 399 | if not cfg.cf_team_domain or not _CF_DOMAIN_RE.match(cfg.cf_team_domain): |
| 400 | raise ValueError(f"Invalid Cloudflare team domain '{cfg.cf_team_domain}'. " "Must match *.cloudflareaccess.com") |
| 401 | |
| 402 | header = jwt.get_unverified_header(token) |
| 403 | kid = header.get("kid") |
| 404 | |
| 405 | keys = await _get_cf_public_keys(cfg.cf_team_domain) |
| 406 | |
| 407 | expected_issuer = f"https://{cfg.cf_team_domain}" |
| 408 | |
| 409 | for k in keys: |
| 410 | if k["kid"] == kid: |
| 411 | claims = jwt.decode( |
| 412 | token, |
| 413 | key=k["key"], |
| 414 | algorithms=["RS256"], |
| 415 | audience=cfg.cf_audience, |
| 416 | issuer=expected_issuer, |
| 417 | options={"verify_exp": True, "verify_aud": True, "verify_iss": True}, |
| 418 | ) |
| 419 | return claims |
| 420 | |
| 421 | raise ValueError("Unable to find matching signing key in Cloudflare JWKS") |