Base

user/user_auth_system/user_management/views/user.py

@@ -0,0 +1,100 @@
+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)