agilentics / boiler
"""Google sign-in: the OAuth 2.0 / OpenID Connect redirect dance.

Two routes, both public (crud.PUBLIC_PATHS) because they run before any session
exists:

  GET /auth/google/start     build Google's consent URL and redirect to it,
                             stashing a `state` (anti-forgery) and `nonce`
                             (anti-replay) in the session first
  GET /auth/google/callback  Google redirects back here with a code; exchange it
                             for an ID token, validate the token's claims,
                             resolve the user (app/core/oauth), start a session

This is the authorization-code flow: the ID token is fetched from Google's token
endpoint directly over TLS, not read from the browser, so its claims can be
trusted without re-verifying the signature - Google's own guidance for this flow.
Issuer, audience, expiry and the nonce we issued are still checked, because those
are what say the token was minted for *this* client and *this* attempt.

The callback's forgery defence is the OAuth `state` parameter rather than our own
CSRF token: there is no session yet to carry one, and `state` is the standard
mechanism for this leg.

Configuration is GOOGLE_CLIENT_ID + GOOGLE_CLIENT_SECRET. Unset means the feature
is off: the routes bounce back to /login and the template hides the button - the
same "unset config = disabled" stance settings.SMTP_HOST takes.
"""
import base64
import json
import logging
import secrets
import time

import httpx
from flask import redirect, request
from flask import session as flask_session

import settings
from app.core import Invalid
from app.core import auth as auth_core
from app.core import oauth as oauth_core
from utils.db_wrapper import db_session

from .crud import public_base_url
from .session import _begin_session

logger = logging.getLogger(__name__)

_AUTH_ENDPOINT = "https://accounts.google.com/o/oauth2/v2/auth"
_TOKEN_ENDPOINT = "https://oauth2.googleapis.com/token"
_VALID_ISSUERS = frozenset({"accounts.google.com", "https://accounts.google.com"})

STATE_KEY = "google_oauth_state"
NONCE_KEY = "google_oauth_nonce"


def is_enabled() -> bool:
    cfg = settings.Config
    return bool(cfg.GOOGLE_CLIENT_ID and cfg.GOOGLE_CLIENT_SECRET)


def _redirect_uri() -> str:
    """The callback URL, which must match one registered in the Google console.
    Built from the canonical public base URL - never from the Host header, which
    an attacker controls."""
    return f"{public_base_url()}/auth/google/callback"


def _fail(reason: str):
    """Every failure lands on the sign-in page with a generic marker. The reason
    goes to the log and not to the query string: it distinguishes "no such
    account" from "bad token", which is not a stranger's business."""
    logger.warning("Google sign-in failed: %s", reason)
    return redirect("/login?error=google")


def page_google_start():
    """Kick off the flow: stash state+nonce, redirect to Google's consent screen.
    If Google sign-in is not configured, quietly return to /login."""
    if not is_enabled():
        return redirect("/login")

    # Fail closed if no canonical base URL is configured on a non-local host,
    # rather than sending Google a redirect_uri derived from the Host header.
    try:
        redirect_uri = _redirect_uri()
    except Invalid as e:
        return _fail(str(e))

    state = secrets.token_urlsafe(24)
    nonce = secrets.token_urlsafe(24)
    flask_session[STATE_KEY] = state
    flask_session[NONCE_KEY] = nonce

    params = {
        "client_id": settings.Config.GOOGLE_CLIENT_ID,
        "redirect_uri": redirect_uri,
        "response_type": "code",
        "scope": "openid email profile",
        "state": state,
        "nonce": nonce,
        # select_account so a shared machine does not silently reuse whichever
        # Google session happens to be open.
        "prompt": "select_account",
    }
    return redirect(str(httpx.URL(_AUTH_ENDPOINT, params=params)))


def _decode_jwt_payload(token: str) -> dict:
    """The claims out of a JWT's payload segment. No signature check: the token
    came from Google's token endpoint over TLS rather than from the browser, so
    its contents are already trusted - the claims are validated below."""
    parts = token.split(".")
    if len(parts) != 3:
        raise ValueError("Malformed ID token.")
    payload_b64 = parts[1]
    padded = payload_b64 + "=" * (-len(payload_b64) % 4)
    return json.loads(base64.urlsafe_b64decode(padded))


def page_google_callback():
    """Finish the flow: verify state, exchange the code, validate the ID token,
    resolve the user, and start a session."""
    if not is_enabled():
        return redirect("/login")

    if request.args.get("error"):
        return _fail(f"consent returned {request.args.get('error')!r}")

    # Popped, not read: state and nonce are single-use, so a replayed callback
    # finds nothing to match against.
    expected_state = flask_session.pop(STATE_KEY, None)
    nonce = flask_session.pop(NONCE_KEY, None)
    if not expected_state or not secrets.compare_digest(
            request.args.get("state", ""), expected_state):
        return _fail("state mismatch")

    code = request.args.get("code")
    if not code:
        return _fail("no code")

    cfg = settings.Config
    try:
        resp = httpx.post(_TOKEN_ENDPOINT, data={
            "code": code,
            "client_id": cfg.GOOGLE_CLIENT_ID,
            "client_secret": cfg.GOOGLE_CLIENT_SECRET,
            "redirect_uri": _redirect_uri(),
            "grant_type": "authorization_code",
        }, timeout=10)
        resp.raise_for_status()
        token_data = resp.json()
    except Exception as e:
        return _fail(f"token exchange failed: {e}")

    id_token = token_data.get("id_token")
    if not id_token:
        return _fail("no id_token in token response")

    try:
        claims = _decode_jwt_payload(id_token)
    except Exception as e:
        return _fail(f"could not decode id_token: {e}")

    if claims.get("iss") not in _VALID_ISSUERS:
        return _fail(f"bad issuer {claims.get('iss')!r}")
    if claims.get("aud") != cfg.GOOGLE_CLIENT_ID:
        return _fail("audience mismatch")
    if nonce and claims.get("nonce") != nonce:
        return _fail("nonce mismatch")
    if claims.get("exp") and int(claims["exp"]) < int(time.time()):
        return _fail("id_token expired")

    # email_verified is the claim that makes case 2 in app/core/oauth safe -
    # without it, an unverified Google address could be used to link, and so to
    # sign in as, an existing account.
    email = claims.get("email")
    if not email or not claims.get("email_verified"):
        return _fail("email missing or unverified")

    try:
        user = oauth_core.find_or_create_google_user(
            db_session, google_sub=claims.get("sub"), email=email,
            name=claims.get("name"),
        )
        resolved = auth_core.principal_for_user(db_session, user)
    except Exception as e:
        db_session.rollback()
        return _fail(f"could not resolve user: {e}")

    _begin_session(user.id, resolved.org_id)
    return redirect("/")