488 lines
16 KiB
Python
488 lines
16 KiB
Python
import json
|
|
import secrets
|
|
|
|
import requests
|
|
from django.conf import settings
|
|
from django.contrib.auth.hashers import make_password
|
|
from drf_spectacular.types import OpenApiTypes
|
|
from drf_spectacular.utils import OpenApiParameter, extend_schema, inline_serializer
|
|
|
|
from rest_framework import serializers, status
|
|
from rest_framework.permissions import IsAuthenticated
|
|
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 common import serializer
|
|
from common.models import Org, Profile, User
|
|
from common.serializer import OrgAwareRefreshToken
|
|
|
|
|
|
class GoogleOAuthCallbackView(APIView):
|
|
"""
|
|
Handle Google OAuth authorization code exchange with PKCE.
|
|
"""
|
|
|
|
permission_classes = []
|
|
authentication_classes = []
|
|
|
|
@extend_schema(
|
|
tags=["auth"],
|
|
request=inline_serializer(
|
|
name="GoogleOAuthRequest",
|
|
fields={
|
|
"code": serializers.CharField(),
|
|
"code_verifier": serializers.CharField(),
|
|
"redirect_uri": serializers.CharField(),
|
|
},
|
|
),
|
|
responses={
|
|
200: inline_serializer(
|
|
name="GoogleOAuthResponse",
|
|
fields={
|
|
"access_token": serializers.CharField(),
|
|
"refresh_token": serializers.CharField(),
|
|
"user": serializers.DictField(),
|
|
},
|
|
)
|
|
},
|
|
)
|
|
def post(self, request):
|
|
import base64
|
|
|
|
from django.conf import settings as django_settings
|
|
from django.utils import timezone
|
|
|
|
code = request.data.get("code")
|
|
code_verifier = request.data.get("code_verifier")
|
|
redirect_uri = request.data.get("redirect_uri")
|
|
|
|
if not all([code, code_verifier, redirect_uri]):
|
|
return Response(
|
|
{"error": "Missing required parameters"},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
|
|
# Exchange code for tokens with Google
|
|
try:
|
|
token_response = requests.post(
|
|
"https://oauth2.googleapis.com/token",
|
|
data={
|
|
"code": code,
|
|
"client_id": django_settings.GOOGLE_CLIENT_ID,
|
|
"client_secret": django_settings.GOOGLE_CLIENT_SECRET,
|
|
"redirect_uri": redirect_uri,
|
|
"grant_type": "authorization_code",
|
|
"code_verifier": code_verifier,
|
|
},
|
|
timeout=30,
|
|
)
|
|
except requests.RequestException:
|
|
return Response(
|
|
{"error": "Failed to communicate with Google"},
|
|
status=status.HTTP_502_BAD_GATEWAY,
|
|
)
|
|
|
|
if token_response.status_code != 200:
|
|
error_data = token_response.json() if token_response.content else {}
|
|
return Response(
|
|
{"error": error_data.get("error_description", "Token exchange failed")},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
|
|
token_data = token_response.json()
|
|
id_token = token_data.get("id_token")
|
|
if not id_token:
|
|
return Response(
|
|
{"error": "No ID token received"},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
|
|
# Decode ID token payload (no verification needed - we just got it from Google over HTTPS)
|
|
try:
|
|
payload_part = id_token.split(".")[1]
|
|
# Add padding if needed
|
|
payload_part += "=" * (4 - len(payload_part) % 4)
|
|
payload = json.loads(base64.urlsafe_b64decode(payload_part))
|
|
email = payload.get("email")
|
|
picture = payload.get("picture", "")
|
|
except (IndexError, ValueError, json.JSONDecodeError):
|
|
return Response(
|
|
{"error": "Invalid ID token format"},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
|
|
if not email:
|
|
return Response(
|
|
{"error": "No email in token"},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
|
|
# Get or create user
|
|
try:
|
|
user = User.objects.get(email=email)
|
|
except User.DoesNotExist:
|
|
user = User.objects.create(
|
|
email=email,
|
|
profile_pic=picture,
|
|
password=make_password(secrets.token_urlsafe(32)),
|
|
)
|
|
|
|
user.last_login = timezone.now()
|
|
user.save(update_fields=["last_login"])
|
|
|
|
# Generate JWT tokens
|
|
token = RefreshToken.for_user(user)
|
|
|
|
return Response(
|
|
{
|
|
"access_token": str(token.access_token),
|
|
"refresh_token": str(token),
|
|
"user": {"id": str(user.id), "email": user.email},
|
|
}
|
|
)
|
|
|
|
|
|
class LoginView(APIView):
|
|
"""
|
|
Login with email and password, returns JWT tokens
|
|
"""
|
|
|
|
permission_classes = []
|
|
authentication_classes = []
|
|
|
|
@extend_schema(
|
|
description="Login with email and password",
|
|
request=serializer.LoginSerializer,
|
|
responses={
|
|
200: {
|
|
"type": "object",
|
|
"properties": {
|
|
"access_token": {"type": "string"},
|
|
"refresh_token": {"type": "string"},
|
|
"user": {
|
|
"type": "object",
|
|
"properties": {
|
|
"id": {"type": "string"},
|
|
"email": {"type": "string"},
|
|
"organizations": {"type": "array"},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
},
|
|
)
|
|
def post(self, request):
|
|
from django.utils import timezone
|
|
|
|
from common.audit_log import audit_log
|
|
|
|
serializer_obj = serializer.LoginSerializer(data=request.data)
|
|
if serializer_obj.is_valid():
|
|
user = serializer_obj.validated_data["user"]
|
|
|
|
# Update last_login timestamp
|
|
user.last_login = timezone.now()
|
|
user.save(update_fields=["last_login"])
|
|
|
|
# Get user's organizations
|
|
profiles = Profile.objects.filter(user=user, is_active=True)
|
|
|
|
# Get the default org (first org or from request)
|
|
org_id = request.data.get("org_id")
|
|
default_org = None
|
|
|
|
if org_id:
|
|
# Use specified org if user has access
|
|
try:
|
|
profile = profiles.get(org_id=org_id)
|
|
default_org = profile.org
|
|
except Profile.DoesNotExist:
|
|
audit_log.login_failure(
|
|
user.email, f"No access to org {org_id}", request
|
|
)
|
|
return Response(
|
|
{
|
|
"error": "User does not have access to specified organization"
|
|
},
|
|
status=status.HTTP_403_FORBIDDEN,
|
|
)
|
|
elif profiles.exists():
|
|
# Use first available org
|
|
profile = profiles.first()
|
|
default_org = profile.org
|
|
|
|
# Generate JWT tokens with org context (include profile for role)
|
|
if default_org:
|
|
token = OrgAwareRefreshToken.for_user_and_org(
|
|
user, default_org, profile
|
|
)
|
|
else:
|
|
# User has no orgs - generate token without org context
|
|
token = RefreshToken.for_user(user)
|
|
|
|
# Audit log successful login
|
|
audit_log.login_success(user, default_org, request)
|
|
|
|
# Get user details with organizations
|
|
user_serializer = serializer.UserDetailSerializer(user)
|
|
|
|
response_data = {
|
|
"access_token": str(token.access_token),
|
|
"refresh_token": str(token),
|
|
"user": user_serializer.data,
|
|
}
|
|
|
|
# Include current org info if available
|
|
if default_org:
|
|
response_data["current_org"] = {
|
|
"id": str(default_org.id),
|
|
"name": default_org.name,
|
|
}
|
|
|
|
return Response(response_data, status=status.HTTP_200_OK)
|
|
|
|
# Log failed login
|
|
email = request.data.get("email", "unknown")
|
|
audit_log.login_failure(email, str(serializer_obj.errors), request)
|
|
return Response(serializer_obj.errors, status=status.HTTP_400_BAD_REQUEST)
|
|
|
|
|
|
class RegisterView(APIView):
|
|
"""
|
|
Register a new user account
|
|
"""
|
|
|
|
permission_classes = []
|
|
authentication_classes = []
|
|
|
|
@extend_schema(
|
|
description="Register a new user account",
|
|
request=serializer.RegisterSerializer,
|
|
responses={
|
|
201: {
|
|
"type": "object",
|
|
"properties": {
|
|
"message": {"type": "string"},
|
|
"user_id": {"type": "string"},
|
|
"email": {"type": "string"},
|
|
},
|
|
}
|
|
},
|
|
)
|
|
def post(self, request):
|
|
serializer_obj = serializer.RegisterSerializer(data=request.data)
|
|
if serializer_obj.is_valid():
|
|
user = serializer_obj.save()
|
|
return Response(
|
|
{
|
|
"message": "User registered successfully. Please check your email for activation link.",
|
|
"user_id": str(user.id),
|
|
"email": user.email,
|
|
},
|
|
status=status.HTTP_201_CREATED,
|
|
)
|
|
|
|
return Response(serializer_obj.errors, status=status.HTTP_400_BAD_REQUEST)
|
|
|
|
|
|
class MeView(APIView):
|
|
"""
|
|
Get current authenticated user details
|
|
"""
|
|
|
|
authentication_classes = [JWTAuthentication]
|
|
permission_classes = [IsAuthenticated]
|
|
|
|
@extend_schema(
|
|
description="Get current authenticated user with organizations",
|
|
responses={200: serializer.UserDetailSerializer},
|
|
)
|
|
def get(self, request):
|
|
user_serializer = serializer.UserDetailSerializer(request.user)
|
|
return Response(user_serializer.data, status=status.HTTP_200_OK)
|
|
|
|
|
|
class OrgAwareTokenRefreshView(APIView):
|
|
"""
|
|
Custom token refresh that validates org membership.
|
|
|
|
When refreshing a token, this view:
|
|
1. Validates the refresh token
|
|
2. Checks that user still has access to the org in the token
|
|
3. Issues new tokens with the same org context
|
|
|
|
If membership was revoked, returns 403 and user must login again.
|
|
"""
|
|
|
|
permission_classes = []
|
|
authentication_classes = []
|
|
|
|
@extend_schema(
|
|
description="Refresh access token with org membership validation",
|
|
request=inline_serializer(
|
|
name="OrgAwareTokenRefreshRequest",
|
|
fields={"refresh": serializers.CharField(help_text="Refresh token")},
|
|
),
|
|
responses={
|
|
200: inline_serializer(
|
|
name="OrgAwareTokenRefreshResponse",
|
|
fields={
|
|
"access": serializers.CharField(),
|
|
"refresh": serializers.CharField(),
|
|
},
|
|
)
|
|
},
|
|
)
|
|
def post(self, request):
|
|
from rest_framework_simplejwt.exceptions import TokenError
|
|
from rest_framework_simplejwt.tokens import RefreshToken as BaseRefreshToken
|
|
|
|
from common.audit_log import audit_log
|
|
|
|
refresh_token = request.data.get("refresh")
|
|
|
|
if not refresh_token:
|
|
return Response(
|
|
{"error": "Refresh token is required"},
|
|
status=status.HTTP_400_BAD_REQUEST,
|
|
)
|
|
|
|
try:
|
|
# Decode and validate refresh token
|
|
token = BaseRefreshToken(refresh_token)
|
|
user_id = token["user_id"]
|
|
org_id = token.get("org_id")
|
|
|
|
# Get user
|
|
user = User.objects.get(id=user_id)
|
|
|
|
if not user.is_active:
|
|
return Response(
|
|
{"error": "User account is disabled"},
|
|
status=status.HTTP_403_FORBIDDEN,
|
|
)
|
|
|
|
# If token has org context, validate membership
|
|
if org_id:
|
|
try:
|
|
profile = Profile.objects.get(
|
|
user=user, org_id=org_id, is_active=True
|
|
)
|
|
org = profile.org
|
|
except Profile.DoesNotExist:
|
|
# Membership revoked - user must login again
|
|
audit_log.token_revoked(
|
|
user, None, f"Membership revoked for org {org_id}", request
|
|
)
|
|
return Response(
|
|
{
|
|
"error": "Organization membership revoked. Please login again."
|
|
},
|
|
status=status.HTTP_403_FORBIDDEN,
|
|
)
|
|
|
|
# Generate new tokens with same org context (include profile for role)
|
|
new_token = OrgAwareRefreshToken.for_user_and_org(user, org, profile)
|
|
audit_log.token_refresh(user, org, request)
|
|
else:
|
|
# No org context - standard refresh
|
|
new_token = BaseRefreshToken.for_user(user)
|
|
audit_log.token_refresh(user, None, request)
|
|
|
|
return Response(
|
|
{"access": str(new_token.access_token), "refresh": str(new_token)},
|
|
status=status.HTTP_200_OK,
|
|
)
|
|
|
|
except TokenError:
|
|
return Response(
|
|
{"error": "Invalid or expired token"},
|
|
status=status.HTTP_401_UNAUTHORIZED,
|
|
)
|
|
except User.DoesNotExist:
|
|
return Response(
|
|
{"error": "User not found"}, status=status.HTTP_401_UNAUTHORIZED
|
|
)
|
|
|
|
|
|
class OrgSwitchView(APIView):
|
|
"""
|
|
Switch to a different organization and get new JWT tokens.
|
|
|
|
This endpoint validates that the user has access to the target
|
|
organization and issues new tokens with the org_id claim.
|
|
"""
|
|
|
|
authentication_classes = [JWTAuthentication]
|
|
permission_classes = [IsAuthenticated]
|
|
|
|
@extend_schema(
|
|
description="Switch to a different organization and get new JWT tokens",
|
|
request=inline_serializer(
|
|
name="OrgSwitchRequest",
|
|
fields={
|
|
"org_id": serializers.UUIDField(help_text="Target organization ID")
|
|
},
|
|
),
|
|
responses={
|
|
200: inline_serializer(
|
|
name="OrgSwitchResponse",
|
|
fields={
|
|
"access_token": serializers.CharField(),
|
|
"refresh_token": serializers.CharField(),
|
|
"current_org": serializers.DictField(),
|
|
"profile": serializers.DictField(),
|
|
},
|
|
)
|
|
},
|
|
)
|
|
def post(self, request):
|
|
from common.audit_log import audit_log
|
|
|
|
org_id = request.data.get("org_id")
|
|
|
|
if not org_id:
|
|
return Response(
|
|
{"error": "org_id is required"}, status=status.HTTP_400_BAD_REQUEST
|
|
)
|
|
|
|
# Get current org for audit logging
|
|
from_org = getattr(request, "org", None)
|
|
|
|
# Validate user has access to the target org
|
|
try:
|
|
profile = Profile.objects.get(
|
|
user=request.user, org_id=org_id, is_active=True
|
|
)
|
|
except Profile.DoesNotExist:
|
|
audit_log.permission_denied(
|
|
request.user, from_org, "ORG_SWITCH", f"org:{org_id}", request
|
|
)
|
|
return Response(
|
|
{"error": "User does not have access to this organization"},
|
|
status=status.HTTP_403_FORBIDDEN,
|
|
)
|
|
|
|
# Generate new tokens with the target org (include profile for role)
|
|
token = OrgAwareRefreshToken.for_user_and_org(
|
|
request.user, profile.org, profile
|
|
)
|
|
|
|
# Audit log the org switch
|
|
audit_log.org_switch(request.user, from_org, profile.org, request)
|
|
|
|
return Response(
|
|
{
|
|
"access_token": str(token.access_token),
|
|
"refresh_token": str(token),
|
|
"current_org": {"id": str(profile.org.id), "name": profile.org.name},
|
|
"profile": {
|
|
"id": str(profile.id),
|
|
"role": profile.role,
|
|
"is_organization_admin": profile.is_organization_admin,
|
|
},
|
|
},
|
|
status=status.HTTP_200_OK,
|
|
)
|