chat_backend/llm_be/chat_backend/oauth.py

"""Google / Microsoft OIDC helpers for SSO (#24)."""

from __future__ import annotations

import logging
from dataclasses import dataclass
from datetime import timedelta
from typing import Any
from urllib.parse import urlencode

import httpx
import jwt
from django.conf import settings
from django.core import signing
from django.utils import timezone

from .models import Company, CustomUser, DriveConnection, OAuthIdentity

logger = logging.getLogger(__name__)

STATE_SALT = "chat_backend.oauth.state"
STATE_MAX_AGE_SECONDS = 600

GOOGLE_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"
GOOGLE_TOKEN_URL = "https://oauth2.googleapis.com/token"
GOOGLE_USERINFO_URL = "https://openidconnect.googleapis.com/v1/userinfo"
GOOGLE_DRIVE_READONLY_SCOPE = "https://www.googleapis.com/auth/drive.readonly"

MICROSOFT_AUTH_URL_TMPL = (
    "https://login.microsoftonline.com/{tenant}/oauth2/v2.0/authorize"
)
MICROSOFT_TOKEN_URL_TMPL = (
    "https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token"
)
MICROSOFT_DRIVE_PERSONAL_SCOPE = "Files.Read"
MICROSOFT_DRIVE_COMPANY_SCOPE = "Files.Read.All Sites.Read.All"

# OAuth intents (#24 login/signup; #47 Drive linking).
LOGIN_INTENTS = {"login", "signup"}
DRIVE_LINK_INTENTS = {"link_drive", "link_company_drive"}
VALID_INTENTS = LOGIN_INTENTS | DRIVE_LINK_INTENTS


class OAuthError(Exception):
    """User-facing OAuth failure with a stable error code for the FE."""

    def __init__(self, code: str, message: str = ""):
        self.code = code
        self.message = message or code
        super().__init__(self.message)


@dataclass(frozen=True)
class ProviderProfile:
    provider: str
    subject: str
    email: str
    email_verified: bool
    first_name: str
    last_name: str
    access_token: str
    refresh_token: str
    expires_in: int | None
    scopes: str
    raw: dict[str, Any]


def provider_configured(provider: str) -> bool:
    if provider == OAuthIdentity.Provider.GOOGLE:
        return bool(settings.GOOGLE_OAUTH_CLIENT_ID and settings.GOOGLE_OAUTH_CLIENT_SECRET)
    if provider == OAuthIdentity.Provider.MICROSOFT:
        return bool(
            settings.MICROSOFT_OAUTH_CLIENT_ID and settings.MICROSOFT_OAUTH_CLIENT_SECRET
        )
    return False


def configured_providers() -> dict[str, bool]:
    return {
        OAuthIdentity.Provider.GOOGLE: provider_configured(OAuthIdentity.Provider.GOOGLE),
        OAuthIdentity.Provider.MICROSOFT: provider_configured(
            OAuthIdentity.Provider.MICROSOFT
        ),
    }


def dump_oauth_state(
    *, provider: str, intent: str, user_id: int | None = None
) -> str:
    payload: dict[str, Any] = {"provider": provider, "intent": intent}
    if user_id is not None:
        payload["user_id"] = user_id
    return signing.dumps(payload, salt=STATE_SALT)


def load_oauth_state(state: str) -> dict[str, Any]:
    try:
        data = signing.loads(state, salt=STATE_SALT, max_age=STATE_MAX_AGE_SECONDS)
    except signing.BadSignature as exc:
        raise OAuthError("invalid_state", "OAuth state is invalid or expired.") from exc
    provider = data.get("provider")
    intent = data.get("intent") or "login"
    if provider not in OAuthIdentity.Provider.values:
        raise OAuthError("invalid_state", "Unknown OAuth provider in state.")
    if intent not in VALID_INTENTS:
        raise OAuthError("invalid_state", "Invalid OAuth intent.")
    result: dict[str, Any] = {"provider": provider, "intent": intent}
    if intent in DRIVE_LINK_INTENTS:
        user_id = data.get("user_id")
        if not user_id:
            raise OAuthError(
                "invalid_state", "OAuth state is missing the linking user."
            )
        result["user_id"] = user_id
    return result


def _microsoft_tenant() -> str:
    return settings.MICROSOFT_OAUTH_TENANT or "common"


def build_authorization_url(
    *, provider: str, redirect_uri: str, state: str, intent: str = "login"
) -> str:
    if not provider_configured(provider):
        raise OAuthError("provider_not_configured", f"{provider} OAuth is not configured.")

    if provider == OAuthIdentity.Provider.GOOGLE:
        scope = "openid email profile"
        if intent in DRIVE_LINK_INTENTS:
            scope = f"{scope} {GOOGLE_DRIVE_READONLY_SCOPE}"
        params = {
            "client_id": settings.GOOGLE_OAUTH_CLIENT_ID,
            "redirect_uri": redirect_uri,
            "response_type": "code",
            "scope": scope,
            "state": state,
            "access_type": "offline",
            "prompt": "select_account consent",
            "include_granted_scopes": "true",
        }
        return f"{GOOGLE_AUTH_URL}?{urlencode(params)}"

    if provider == OAuthIdentity.Provider.MICROSOFT:
        scope = "openid email profile offline_access"
        if intent == "link_drive":
            scope = f"{scope} {MICROSOFT_DRIVE_PERSONAL_SCOPE}"
        elif intent == "link_company_drive":
            scope = f"{scope} {MICROSOFT_DRIVE_COMPANY_SCOPE}"
        params = {
            "client_id": settings.MICROSOFT_OAUTH_CLIENT_ID,
            "redirect_uri": redirect_uri,
            "response_type": "code",
            "response_mode": "query",
            "scope": scope,
            "state": state,
            "prompt": "select_account",
        }
        auth_url = MICROSOFT_AUTH_URL_TMPL.format(tenant=_microsoft_tenant())
        return f"{auth_url}?{urlencode(params)}"

    raise OAuthError("invalid_provider", f"Unsupported provider: {provider}")


def _decode_id_token_claims(id_token: str | None) -> dict[str, Any]:
    if not id_token:
        return {}
    # Signature verified via TLS token endpoint + client secret; claims are trusted.
    return jwt.decode(
        id_token,
        options={"verify_signature": False, "verify_aud": False},
    )


def exchange_code_for_profile(
    *, provider: str, code: str, redirect_uri: str
) -> ProviderProfile:
    if provider == OAuthIdentity.Provider.GOOGLE:
        return _exchange_google(code=code, redirect_uri=redirect_uri)
    if provider == OAuthIdentity.Provider.MICROSOFT:
        return _exchange_microsoft(code=code, redirect_uri=redirect_uri)
    raise OAuthError("invalid_provider", f"Unsupported provider: {provider}")


def _exchange_google(*, code: str, redirect_uri: str) -> ProviderProfile:
    with httpx.Client(timeout=20.0) as client:
        token_response = client.post(
            GOOGLE_TOKEN_URL,
            data={
                "code": code,
                "client_id": settings.GOOGLE_OAUTH_CLIENT_ID,
                "client_secret": settings.GOOGLE_OAUTH_CLIENT_SECRET,
                "redirect_uri": redirect_uri,
                "grant_type": "authorization_code",
            },
        )
        if token_response.status_code >= 400:
            logger.warning("Google token exchange failed: %s", token_response.text)
            raise OAuthError("token_exchange_failed", "Google token exchange failed.")
        token_data = token_response.json()
        access_token = token_data.get("access_token") or ""
        if not access_token:
            raise OAuthError("token_exchange_failed", "Google did not return an access token.")

        userinfo_response = client.get(
            GOOGLE_USERINFO_URL,
            headers={"Authorization": f"Bearer {access_token}"},
        )
        if userinfo_response.status_code >= 400:
            logger.warning("Google userinfo failed: %s", userinfo_response.text)
            raise OAuthError("profile_fetch_failed", "Could not load Google profile.")
        profile = userinfo_response.json()

    claims = _decode_id_token_claims(token_data.get("id_token"))
    email = (profile.get("email") or claims.get("email") or "").strip().lower()
    email_verified = bool(
        profile.get("email_verified", claims.get("email_verified", False))
    )
    subject = str(profile.get("sub") or claims.get("sub") or "").strip()
    if not subject:
        raise OAuthError("profile_incomplete", "Google profile missing subject.")

    return ProviderProfile(
        provider=OAuthIdentity.Provider.GOOGLE,
        subject=subject,
        email=email,
        email_verified=email_verified,
        first_name=(profile.get("given_name") or claims.get("given_name") or "").strip(),
        last_name=(profile.get("family_name") or claims.get("family_name") or "").strip(),
        access_token=access_token,
        refresh_token=token_data.get("refresh_token") or "",
        expires_in=_as_int(token_data.get("expires_in")),
        scopes=token_data.get("scope") or "openid email profile",
        raw={"userinfo": profile, "id_token_claims": claims},
    )


def _exchange_microsoft(*, code: str, redirect_uri: str) -> ProviderProfile:
    token_url = MICROSOFT_TOKEN_URL_TMPL.format(tenant=_microsoft_tenant())
    with httpx.Client(timeout=20.0) as client:
        token_response = client.post(
            token_url,
            data={
                "code": code,
                "client_id": settings.MICROSOFT_OAUTH_CLIENT_ID,
                "client_secret": settings.MICROSOFT_OAUTH_CLIENT_SECRET,
                "redirect_uri": redirect_uri,
                "grant_type": "authorization_code",
                "scope": "openid email profile offline_access",
            },
        )
        if token_response.status_code >= 400:
            logger.warning("Microsoft token exchange failed: %s", token_response.text)
            raise OAuthError("token_exchange_failed", "Microsoft token exchange failed.")
        token_data = token_response.json()

    claims = _decode_id_token_claims(token_data.get("id_token"))
    email = (
        claims.get("email")
        or claims.get("preferred_username")
        or claims.get("upn")
        or ""
    )
    email = str(email).strip().lower()
    # Microsoft issues verified tenant emails; treat presence as verified when claim missing.
    email_verified = bool(claims.get("email_verified", True if email else False))
    subject = str(claims.get("oid") or claims.get("sub") or "").strip()
    if not subject:
        raise OAuthError("profile_incomplete", "Microsoft profile missing subject.")

    name = (claims.get("name") or "").strip()
    first_name = (claims.get("given_name") or "").strip()
    last_name = (claims.get("family_name") or "").strip()
    if not first_name and name:
        parts = name.split(" ", 1)
        first_name = parts[0]
        last_name = parts[1] if len(parts) > 1 else ""

    return ProviderProfile(
        provider=OAuthIdentity.Provider.MICROSOFT,
        subject=subject,
        email=email,
        email_verified=email_verified,
        first_name=first_name,
        last_name=last_name,
        access_token=token_data.get("access_token") or "",
        refresh_token=token_data.get("refresh_token") or "",
        expires_in=_as_int(token_data.get("expires_in")),
        scopes=token_data.get("scope") or "openid email profile offline_access",
        raw={"id_token_claims": claims},
    )


def _as_int(value: Any) -> int | None:
    try:
        return int(value) if value is not None else None
    except (TypeError, ValueError):
        return None


def _token_expiry(expires_in: int | None):
    if not expires_in:
        return None
    return timezone.now() + timedelta(seconds=expires_in)


def upsert_identity(user: CustomUser, profile: ProviderProfile) -> OAuthIdentity:
    identity = OAuthIdentity.objects.filter(
        provider=profile.provider, subject=profile.subject
    ).first()
    expires_at = _token_expiry(profile.expires_in)
    if identity:
        identity.user = user
        identity.email = profile.email
        identity.access_token = profile.access_token
        if profile.refresh_token:
            identity.refresh_token = profile.refresh_token
        identity.token_expires_at = expires_at
        identity.scopes = profile.scopes
        identity.raw_profile = profile.raw
        identity.save()
        return identity

    return OAuthIdentity.objects.create(
        user=user,
        provider=profile.provider,
        subject=profile.subject,
        email=profile.email,
        access_token=profile.access_token,
        refresh_token=profile.refresh_token or "",
        token_expires_at=expires_at,
        scopes=profile.scopes,
        raw_profile=profile.raw,
    )


def upsert_drive_connection(
    *, user: CustomUser, kind: str, profile: ProviderProfile
) -> DriveConnection:
    """Create/refresh a DriveConnection from a link_drive/link_company_drive callback (#47).

    Personal connections do not require a company (#55). Company connections do.
    """
    if kind == DriveConnection.Kind.COMPANY and not user.company_id:
        raise OAuthError(
            "no_company",
            "A company is required before connecting a company Drive.",
        )

    expires_at = _token_expiry(profile.expires_in)
    if kind == DriveConnection.Kind.PERSONAL:
        connection = DriveConnection.objects.filter(
            provider=profile.provider,
            kind=DriveConnection.Kind.PERSONAL,
            user=user,
        ).first()
        if connection is None:
            connection = DriveConnection(
                company=user.company,  # may be None for solo users
                provider=profile.provider,
                kind=DriveConnection.Kind.PERSONAL,
                user=user,
            )
    else:
        connection = DriveConnection.objects.filter(
            company=user.company,
            provider=profile.provider,
            kind=DriveConnection.Kind.COMPANY,
            user=None,
        ).first()
        if connection is None:
            connection = DriveConnection(
                company=user.company,
                provider=profile.provider,
                kind=DriveConnection.Kind.COMPANY,
                user=None,
            )

    connection.access_token = profile.access_token
    if profile.refresh_token:
        connection.refresh_token = profile.refresh_token
    connection.token_expires_at = expires_at
    connection.scopes = profile.scopes
    connection.external_account_email = profile.email
    connection.is_active = True
    connection.last_sync_error = ""
    connection.save()
    return connection


def _create_sso_user(profile: ProviderProfile) -> CustomUser:
    from monetization.services.plans import try_redeem_backer_email

    company = Company.objects.create(
        name=f"{profile.email}'s workspace",
        state="NA",
        zipcode="00000",
        address="N/A",
    )
    user = CustomUser(
        username=profile.email,
        email=profile.email,
        first_name=profile.first_name,
        last_name=profile.last_name,
        company=company,
        is_company_manager=True,
    )
    user.set_unusable_password()
    user.save()
    try_redeem_backer_email(user)
    return user


def resolve_link_user(user_id: int) -> CustomUser:
    """Load the authenticated user a Drive-link callback should attach to."""
    user = CustomUser.objects.filter(pk=user_id, deleted=False).first()
    if user is None:
        raise OAuthError("user_not_found", "Linking user account was not found.")
    return user


def resolve_user_from_profile(
    *, profile: ProviderProfile, intent: str
) -> tuple[CustomUser, bool]:
    """
    Map IdP profile → CustomUser.

    Returns (user, created).
    """
    if not profile.email:
        raise OAuthError("email_missing", "Email was not provided by the identity provider.")
    if not profile.email_verified:
        raise OAuthError("email_unverified", "Email from the identity provider is not verified.")

    existing_identity = (
        OAuthIdentity.objects.select_related("user")
        .filter(provider=profile.provider, subject=profile.subject)
        .first()
    )
    if existing_identity:
        return existing_identity.user, False

    email_user = (
        CustomUser.objects.filter(email__iexact=profile.email).first()
        or CustomUser.objects.filter(username__iexact=profile.email).first()
    )
    if email_user:
        # Same provider already linked to a different subject → unsafe collision.
        other = (
            OAuthIdentity.objects.filter(provider=profile.provider, user=email_user)
            .exclude(subject=profile.subject)
            .first()
        )
        if other:
            raise OAuthError(
                "link_conflict",
                "This email is already linked to a different identity for this provider.",
            )
        upsert_identity(email_user, profile)
        return email_user, False

    # New account path. Password registration does not enable SSO sign-up.
    allow_create = settings.ENABLE_ACCOUNT_REGISTRATION and settings.ENABLE_SSO
    if intent == "signup" and not allow_create:
        raise OAuthError("registration_disabled", "Account registration is disabled.")
    if intent == "login" and not allow_create:
        raise OAuthError(
            "account_not_found",
            "No account exists for this email. Contact your administrator.",
        )
    if not allow_create:
        raise OAuthError("registration_disabled", "Account registration is disabled.")

    user = _create_sso_user(profile)
    upsert_identity(user, profile)
    return user, True