100 lines
3.5 KiB
Python
100 lines
3.5 KiB
Python
from rest_framework import generics, permissions, status
|
|
from rest_framework.response import Response
|
|
from rest_framework.views import APIView
|
|
from ..models import CustomUser
|
|
from ..serializers import UserSerializer
|
|
from django.contrib.auth import authenticate
|
|
from rest_framework_simplejwt.tokens import RefreshToken
|
|
import pyotp
|
|
import qrcode
|
|
import base64
|
|
from io import BytesIO
|
|
|
|
class UserRegistrationView(generics.CreateAPIView):
|
|
queryset = CustomUser.objects.all()
|
|
serializer_class = UserSerializer
|
|
permission_classes = [permissions.AllowAny]
|
|
|
|
class UserLoginView(APIView):
|
|
permission_classes = [permissions.AllowAny]
|
|
|
|
def post(self, request):
|
|
username = request.data.get('username')
|
|
password = request.data.get('password')
|
|
user = authenticate(username=username, password=password)
|
|
|
|
if user:
|
|
if user.is_2fa_enabled:
|
|
return Response({"message": "2FA is enabled. Please provide OTP.", "require_2fa": True, "user_id": user.id})
|
|
refresh = RefreshToken.for_user(user)
|
|
return Response({
|
|
'refresh': str(refresh),
|
|
'access': str(refresh.access_token),
|
|
'user': UserSerializer(user).data
|
|
})
|
|
return Response({"error": "Invalid Credentials"}, status=status.HTTP_400_BAD_REQUEST)
|
|
|
|
class UserLogoutView(APIView):
|
|
permission_classes = [permissions.IsAuthenticated]
|
|
|
|
def post(self, request):
|
|
try:
|
|
refresh_token = request.data["refresh_token"]
|
|
token = RefreshToken(refresh_token)
|
|
token.blacklist()
|
|
return Response({"message": "Successfully Logged out."}, status=status.HTTP_200_OK)
|
|
except Exception as e:
|
|
return Response({"error": str(e)}, status=status.HTTP_400_BAD_REQUEST)
|
|
|
|
class UserProfileView(generics.RetrieveUpdateAPIView):
|
|
serializer_class = UserSerializer
|
|
permission_classes = [permissions.IsAuthenticated]
|
|
|
|
def get_object(self):
|
|
return self.request.user
|
|
|
|
class Enable2FAView(APIView):
|
|
permission_classes = [permissions.IsAuthenticated]
|
|
|
|
def post(self, request):
|
|
user = request.user
|
|
secret_key = pyotp.random_base32()
|
|
user.otp_secret = secret_key
|
|
user.save()
|
|
|
|
totp = pyotp.TOTP(secret_key)
|
|
uri = totp.provisioning_uri(name=user.email, issuer_name="ArchiveApp")
|
|
|
|
qr = qrcode.QRCode(version=1, box_size=10, border=5)
|
|
qr.add_data(uri)
|
|
qr.make(fit=True)
|
|
img = qr.make_image(fill_color="black", back_color="white")
|
|
|
|
buffered = BytesIO()
|
|
img.save(buffered, format="PNG")
|
|
img_str = base64.b64encode(buffered.getvalue()).decode()
|
|
|
|
return Response({
|
|
'secret_key': secret_key,
|
|
'qr_code': f"data:image/png;base64,{img_str}"
|
|
})
|
|
|
|
class Verify2FAView(APIView):
|
|
permission_classes = [permissions.IsAuthenticated]
|
|
|
|
def post(self, request):
|
|
user = request.user
|
|
otp = request.data.get('otp')
|
|
totp = pyotp.TOTP(user.otp_secret)
|
|
|
|
if totp.verify(otp):
|
|
user.is_2fa_enabled = True
|
|
user.save()
|
|
refresh = RefreshToken.for_user(user)
|
|
return Response({
|
|
"message": "2FA verified successfully",
|
|
'refresh': str(refresh),
|
|
'access': str(refresh.access_token),
|
|
'user': UserSerializer(user).data
|
|
})
|
|
return Response({"error": "Invalid OTP"}, status=status.HTTP_400_BAD_REQUEST) |