westfarn
0d6bf2b024
Personal Drive connections and document workspaces are user-owned and no longer require company_id. Company Drive still requires a company manager. Chat/document scope falls back to a personal workspace for solo users instead of rejecting with company_missing.
266 lines
10 KiB
Python
266 lines
10 KiB
Python
"""OAuth SSO start + callback views (#24)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from urllib.parse import urlencode
|
|
|
|
from django.conf import settings
|
|
from django.http import HttpResponseRedirect
|
|
from django.urls import reverse
|
|
from rest_framework import permissions, status
|
|
from rest_framework.response import Response
|
|
from rest_framework.views import APIView
|
|
from rest_framework_simplejwt.authentication import JWTAuthentication
|
|
from rest_framework_simplejwt.tokens import RefreshToken
|
|
|
|
from finance.services.quotas import FeatureNotAllowed, assert_feature_allowed
|
|
|
|
from .models import DriveConnection, OAuthIdentity
|
|
from .oauth import (
|
|
DRIVE_LINK_INTENTS,
|
|
OAuthError,
|
|
build_authorization_url,
|
|
configured_providers,
|
|
dump_oauth_state,
|
|
exchange_code_for_profile,
|
|
load_oauth_state,
|
|
provider_configured,
|
|
resolve_link_user,
|
|
resolve_user_from_profile,
|
|
upsert_drive_connection,
|
|
upsert_identity,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _callback_redirect_uri(request, provider: str) -> str:
|
|
"""Absolute backend callback URL registered with the IdP."""
|
|
path = reverse("oauth_callback", kwargs={"provider": provider})
|
|
base = (settings.OAUTH_CALLBACK_BASE_URL or "").rstrip("/")
|
|
if base:
|
|
return f"{base}{path}"
|
|
return request.build_absolute_uri(path)
|
|
|
|
|
|
def _frontend_callback_url(**params: str) -> str:
|
|
base = settings.FRONTEND_BASE_URL.rstrip("/")
|
|
query = urlencode({k: v for k, v in params.items() if v is not None and v != ""})
|
|
return f"{base}/auth/callback/?{query}"
|
|
|
|
|
|
def _frontend_drive_redirect_url(**params: str) -> str:
|
|
"""Drive-link callbacks return to Documents storage (#47 / #83)."""
|
|
base = settings.FRONTEND_BASE_URL.rstrip("/")
|
|
query = urlencode({k: v for k, v in params.items() if v is not None and v != ""})
|
|
return f"{base}/document_storage/?{query}"
|
|
|
|
|
|
def _redirect_error(code: str, message: str = "") -> HttpResponseRedirect:
|
|
return HttpResponseRedirect(
|
|
_frontend_callback_url(error=code, error_description=message or code)
|
|
)
|
|
|
|
|
|
def _redirect_drive_error(code: str, message: str = "") -> HttpResponseRedirect:
|
|
return HttpResponseRedirect(
|
|
_frontend_drive_redirect_url(error=code, error_description=message or code)
|
|
)
|
|
|
|
|
|
class OAuthStartView(APIView):
|
|
"""Redirect the browser to Google / Microsoft authorize URL."""
|
|
|
|
permission_classes = (permissions.AllowAny,)
|
|
# Login/signup are anonymous; link_drive/link_company_drive need
|
|
# request.user, so JWT auth runs but never blocks the anonymous flows.
|
|
authentication_classes = (JWTAuthentication,)
|
|
http_method_names = ["get"]
|
|
|
|
def get(self, request, provider: str):
|
|
provider = (provider or "").lower()
|
|
if provider not in OAuthIdentity.Provider.values:
|
|
return Response(
|
|
{"detail": "Unsupported OAuth provider."},
|
|
status=status.HTTP_404_NOT_FOUND,
|
|
)
|
|
if not provider_configured(provider):
|
|
return Response(
|
|
{"detail": f"{provider} OAuth is not configured."},
|
|
status=status.HTTP_503_SERVICE_UNAVAILABLE,
|
|
)
|
|
|
|
intent = (request.query_params.get("intent") or "login").lower()
|
|
if intent not in {"login", "signup", *DRIVE_LINK_INTENTS}:
|
|
return Response(
|
|
{
|
|
"detail": (
|
|
"intent must be one of 'login', 'signup', "
|
|
"'link_drive', 'link_company_drive'."
|
|
)
|
|
},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
if intent == "signup" and not settings.ENABLE_ACCOUNT_REGISTRATION:
|
|
return Response(
|
|
{"detail": "Account registration is disabled."},
|
|
status=status.HTTP_403_FORBIDDEN,
|
|
)
|
|
|
|
user_id = None
|
|
if intent in DRIVE_LINK_INTENTS:
|
|
user = request.user
|
|
if user is None or not user.is_authenticated:
|
|
return Response(
|
|
{"detail": "Authentication is required to link a Drive account."},
|
|
status=status.HTTP_401_UNAUTHORIZED,
|
|
)
|
|
try:
|
|
assert_feature_allowed(user, "rag")
|
|
except FeatureNotAllowed as exc:
|
|
return Response(
|
|
{"detail": exc.message, "code": exc.code},
|
|
status=status.HTTP_403_FORBIDDEN,
|
|
)
|
|
if intent == "link_company_drive" and not user.is_company_manager:
|
|
return Response(
|
|
{"detail": "Only a company manager can connect a company Drive."},
|
|
status=status.HTTP_403_FORBIDDEN,
|
|
)
|
|
user_id = user.id
|
|
|
|
state = dump_oauth_state(provider=provider, intent=intent, user_id=user_id)
|
|
redirect_uri = _callback_redirect_uri(request, provider)
|
|
try:
|
|
auth_url = build_authorization_url(
|
|
provider=provider, redirect_uri=redirect_uri, state=state, intent=intent
|
|
)
|
|
except OAuthError as exc:
|
|
return Response({"detail": exc.message}, status=status.HTTP_400_BAD_REQUEST)
|
|
|
|
# Authenticated Drive-link flows are started via XHR (Bearer JWT). A
|
|
# full-page redirect would drop the Authorization header, so return
|
|
# the IdP URL as JSON when the client asks for it (?response=json).
|
|
if intent in DRIVE_LINK_INTENTS and request.query_params.get("response") == "json":
|
|
return Response({"authorize_url": auth_url})
|
|
|
|
return HttpResponseRedirect(auth_url)
|
|
|
|
|
|
class OAuthCallbackView(APIView):
|
|
"""IdP redirect target — exchange code, create/link user, send JWTs to FE."""
|
|
|
|
permission_classes = (permissions.AllowAny,)
|
|
authentication_classes = ()
|
|
http_method_names = ["get"]
|
|
|
|
def get(self, request, provider: str):
|
|
provider = (provider or "").lower()
|
|
if provider not in OAuthIdentity.Provider.values:
|
|
return _redirect_error("invalid_provider", "Unsupported OAuth provider.")
|
|
|
|
error = request.query_params.get("error")
|
|
if error:
|
|
description = request.query_params.get("error_description") or error
|
|
code = "access_denied" if error == "access_denied" else "provider_error"
|
|
return _redirect_error(code, description)
|
|
|
|
code = request.query_params.get("code")
|
|
state = request.query_params.get("state")
|
|
if not code or not state:
|
|
return _redirect_error("missing_code", "Missing OAuth code or state.")
|
|
|
|
try:
|
|
state_data = load_oauth_state(state)
|
|
if state_data["provider"] != provider:
|
|
raise OAuthError("invalid_state", "OAuth provider mismatch.")
|
|
except OAuthError as exc:
|
|
logger.info("OAuth callback failed (%s): %s", exc.code, exc.message)
|
|
return _redirect_error(exc.code, exc.message)
|
|
|
|
intent = state_data["intent"]
|
|
if intent in DRIVE_LINK_INTENTS:
|
|
return self._handle_drive_link_callback(
|
|
request, provider=provider, code=code, state_data=state_data
|
|
)
|
|
|
|
try:
|
|
redirect_uri = _callback_redirect_uri(request, provider)
|
|
profile = exchange_code_for_profile(
|
|
provider=provider, code=code, redirect_uri=redirect_uri
|
|
)
|
|
user, created = resolve_user_from_profile(profile=profile, intent=intent)
|
|
# Refresh stored tokens on every successful login.
|
|
upsert_identity(user, profile)
|
|
except OAuthError as exc:
|
|
logger.info("OAuth callback failed (%s): %s", exc.code, exc.message)
|
|
return _redirect_error(exc.code, exc.message)
|
|
except Exception:
|
|
logger.exception("Unexpected OAuth callback failure")
|
|
return _redirect_error("server_error", "Unexpected OAuth error.")
|
|
|
|
refresh = RefreshToken.for_user(user)
|
|
from finance.services.plans import needs_checkout as user_needs_checkout
|
|
|
|
needs_checkout = "1" if (created and user_needs_checkout(user)) else "0"
|
|
return HttpResponseRedirect(
|
|
_frontend_callback_url(
|
|
access=str(refresh.access_token),
|
|
refresh=str(refresh),
|
|
created="1" if created else "0",
|
|
needs_checkout=needs_checkout,
|
|
)
|
|
)
|
|
|
|
def _handle_drive_link_callback(
|
|
self, request, *, provider: str, code: str, state_data: dict
|
|
) -> HttpResponseRedirect:
|
|
"""Exchange code + upsert a DriveConnection for an already-authenticated user (#47)."""
|
|
try:
|
|
user = resolve_link_user(state_data["user_id"])
|
|
kind = (
|
|
DriveConnection.Kind.COMPANY
|
|
if state_data["intent"] == "link_company_drive"
|
|
else DriveConnection.Kind.PERSONAL
|
|
)
|
|
if kind == DriveConnection.Kind.COMPANY and not user.is_company_manager:
|
|
raise OAuthError(
|
|
"forbidden",
|
|
"Only a company manager can connect a company Drive.",
|
|
)
|
|
if kind == DriveConnection.Kind.COMPANY and not user.company_id:
|
|
raise OAuthError(
|
|
"no_company",
|
|
"A company is required before connecting a company Drive.",
|
|
)
|
|
assert_feature_allowed(user, "rag")
|
|
|
|
redirect_uri = _callback_redirect_uri(request, provider)
|
|
profile = exchange_code_for_profile(
|
|
provider=provider, code=code, redirect_uri=redirect_uri
|
|
)
|
|
connection = upsert_drive_connection(user=user, kind=kind, profile=profile)
|
|
except FeatureNotAllowed as exc:
|
|
logger.info("Drive link denied by feature gate (%s): %s", exc.code, exc.message)
|
|
return _redirect_drive_error(exc.code, exc.message)
|
|
except OAuthError as exc:
|
|
logger.info("Drive link callback failed (%s): %s", exc.code, exc.message)
|
|
return _redirect_drive_error(exc.code, exc.message)
|
|
except Exception:
|
|
logger.exception("Unexpected Drive link callback failure")
|
|
return _redirect_drive_error("server_error", "Unexpected Drive link error.")
|
|
|
|
return HttpResponseRedirect(
|
|
_frontend_drive_redirect_url(
|
|
drive_connected="1",
|
|
provider=provider,
|
|
kind=connection.kind,
|
|
)
|
|
)
|
|
|
|
|
|
def oauth_public_flags() -> dict:
|
|
"""Feature flags for /public/settings/."""
|
|
return configured_providers()
|