thermograph/discord_link.py

"""Link a Thermograph account to a Discord account via OAuth2 (identify scope).

The Discord user id we store here is the durable key a later feature (DM alerts)
uses to reach the person. The flow is the standard authorization-code grant:

  /link/start    -> redirect the logged-in user to Discord's consent screen
  /link/callback -> exchange the code, read the Discord user id, store it
  /unlink        -> forget it

`state` is signed with the app's auth secret (stdlib hmac, no new dependency) and
carries the Thermograph user id, so the callback can't be replayed or bound to a
different account. Everything requires an active session, so linking always acts on
the person who is actually logged in.
"""
from __future__ import annotations

import base64
import hashlib
import hmac
import json
import os
import time
import urllib.parse

import httpx
from fastapi import APIRouter, Depends, Request
from fastapi.responses import RedirectResponse
from sqlalchemy import update
from sqlalchemy.ext.asyncio import AsyncSession

from db import get_async_session
from models import User
from users import SECRET, current_active_user

router = APIRouter(tags=["discord"])

CLIENT_ID = os.environ.get("THERMOGRAPH_DISCORD_APP_ID", "").strip()
CLIENT_SECRET = os.environ.get("THERMOGRAPH_DISCORD_CLIENT_SECRET", "").strip()
BASE = os.environ.get("THERMOGRAPH_BASE", "/thermograph").rstrip("/")

_AUTHORIZE = "https://discord.com/api/oauth2/authorize"
_TOKEN = "https://discord.com/api/oauth2/token"
_ME = "https://discord.com/api/users/@me"

# How long a start->callback round trip may take before the signed state expires.
_STATE_TTL = 600
_TIMEOUT = httpx.Timeout(10.0)


def enabled() -> bool:
    return bool(CLIENT_ID and CLIENT_SECRET)


# --- signed state ------------------------------------------------------------
def _b64(raw: bytes) -> str:
    return base64.urlsafe_b64encode(raw).decode().rstrip("=")


def _unb64(s: str) -> bytes:
    return base64.urlsafe_b64decode(s + "=" * (-len(s) % 4))


def _sign_state(uid: str) -> str:
    payload = _b64(json.dumps({"u": uid, "t": int(time.time())}).encode())
    mac = hmac.new(SECRET.encode(), payload.encode(), hashlib.sha256).hexdigest()[:32]
    return f"{payload}.{mac}"


def _verify_state(state: str) -> str | None:
    """The Thermograph user id the state was minted for, or None if it's missing,
    tampered, or older than the TTL."""
    try:
        payload, mac = state.split(".", 1)
    except (ValueError, AttributeError):
        return None
    expect = hmac.new(SECRET.encode(), payload.encode(), hashlib.sha256).hexdigest()[:32]
    if not hmac.compare_digest(mac, expect):
        return None
    try:
        data = json.loads(_unb64(payload))
    except (ValueError, json.JSONDecodeError):
        return None
    if time.time() - float(data.get("t", 0)) > _STATE_TTL:
        return None
    return data.get("u")


def _redirect_uri(request: Request) -> str:
    """The callback URL, matched to how the app is actually reached (scheme/host +
    base path). Must equal the redirect registered in the Discord portal."""
    proto = request.headers.get("x-forwarded-proto") or request.url.scheme
    host = request.headers.get("host") or request.url.netloc
    return f"{proto}://{host}{BASE}/api/v2/discord/link/callback"


def _account_redirect(status: str) -> RedirectResponse:
    # Back to the alerts page, which surfaces the link state; 303 so the browser
    # issues a GET after the callback.
    return RedirectResponse(url=f"{BASE}/subscriptions?discord={status}", status_code=303)


# --- routes ------------------------------------------------------------------
@router.get("/discord/link/start")
async def link_start(request: Request, user=Depends(current_active_user)):
    if not enabled():
        return _account_redirect("unavailable")
    params = {
        "client_id": CLIENT_ID,
        "redirect_uri": _redirect_uri(request),
        "response_type": "code",
        "scope": "identify",
        "state": _sign_state(str(user.id)),
        "prompt": "consent",
    }
    return RedirectResponse(url=f"{_AUTHORIZE}?{urllib.parse.urlencode(params)}", status_code=303)


@router.get("/discord/link/callback")
async def link_callback(
    request: Request,
    user=Depends(current_active_user),
    session: AsyncSession = Depends(get_async_session),
):
    if not enabled():
        return _account_redirect("unavailable")
    # User declined on Discord's screen, or the state doesn't check out.
    code = request.query_params.get("code")
    state = request.query_params.get("state", "")
    if not code:
        return _account_redirect("cancelled")
    if _verify_state(state) != str(user.id):
        return _account_redirect("error")
    try:
        async with httpx.AsyncClient(timeout=_TIMEOUT) as client:
            tok = await client.post(_TOKEN, data={
                "client_id": CLIENT_ID,
                "client_secret": CLIENT_SECRET,
                "grant_type": "authorization_code",
                "code": code,
                "redirect_uri": _redirect_uri(request),
            }, headers={"Content-Type": "application/x-www-form-urlencoded"})
            if tok.status_code >= 300:
                return _account_redirect("error")
            access = tok.json().get("access_token")
            me = await client.get(_ME, headers={"Authorization": f"Bearer {access}"})
            if me.status_code >= 300:
                return _account_redirect("error")
            discord_id = str(me.json().get("id") or "")
    except Exception:  # noqa: BLE001 - any network/parse failure is a graceful error
        return _account_redirect("error")
    if not discord_id:
        return _account_redirect("error")
    await session.execute(
        update(User).where(User.id == user.id).values(discord_id=discord_id)
    )
    await session.commit()
    return _account_redirect("linked")


@router.post("/discord/unlink", status_code=204)
async def unlink(
    user=Depends(current_active_user),
    session: AsyncSession = Depends(get_async_session),
):
    await session.execute(
        update(User).where(User.id == user.id).values(discord_id=None)
    )
    await session.commit()