Merge pull request #130 from DocShow-AI/login_fix

Login fix

backend/databases/user_manager.py

@@ -25,7 +25,19 @@ def __init__(self, session: Session):
         """
         self.db_session = session
 
-    def get_user(self, username: str):
+    def get_user_by_email(self, email: str) -> User:
+        """
+        Get a user based on their email.
+
+        Args:
+            email (str): The email of the user.
+
+        Returns:
+            User: The User object if found, else None.
+        """
+        return self.db_session.query(User).filter(User.email == email).first()
+
+    def get_user_by_username(self, username: str) -> User:
         """
         Get a user based on their username.
 

backend/envs/dev/initialization/setup_dev_environment.py

@@ -41,7 +41,7 @@ def create_admin_user():
 
     with DatabaseManager() as session:
         user_manager = UserManager(session)
-        existing_user = user_manager.get_user(admin_user.username)
+        existing_user = user_manager.get_user_by_username(admin_user.username)
         if not existing_user:
             user_manager.create_user(admin_user)
             logger.debug("Admin user created.")

backend/models/auth.py

@@ -0,0 +1,17 @@
+from pydantic import BaseModel, EmailStr
+from typing import Optional
+
+
+class CustomOAuth2PasswordRequestForm(BaseModel):
+    email: Optional[EmailStr] = None
+    username: Optional[str] = None
+    password: str
+
+    class Config:
+        json_schema_extra = {
+            "example": {
+                "email": "[email protected]",
+                "username": None,
+                "password": "secret",
+            }
+        }

backend/routes/auth_routes.py

@@ -4,12 +4,13 @@
 the security module for creating JWT tokens and password hashing/verification.
 """
 from datetime import timedelta
-from fastapi import APIRouter, Depends, HTTPException, Response, status
-from fastapi.security import OAuth2PasswordRequestForm
-from pydantic import BaseModel
+from fastapi import APIRouter, Depends, Form, HTTPException, Response, status
+from pydantic import BaseModel, EmailStr
+from typing import Optional
 
 from databases.database_manager import DatabaseManager
 from databases.user_manager import UserManager
+from models.auth import CustomOAuth2PasswordRequestForm
 from models.user import User, UserCreate
 from security import (
     authenticate_user,
@@ -40,7 +41,10 @@ class LogoutResponse(BaseModel):
 
 @auth_router.post("/token/", response_model=LoginResponse)
 async def login_for_access_token(
-    response: Response, form_data: OAuth2PasswordRequestForm = Depends()
+    response: Response,
+    username: Optional[str] = Form(None),
+    email: Optional[EmailStr] = Form(None),
+    password: str = Form(...),
 ):
     """
     Authenticate a user and set a JWT token in a cookie upon successful authentication.
@@ -49,12 +53,16 @@ async def login_for_access_token(
     it creates a JWT token and sets it in a secure, HttpOnly cookie in the response.
 
     Args:
-        form_data (OAuth2PasswordRequestForm): Contains the user's provided username and password.
+        form_data (CustomOAuth2PasswordRequestForm): Contains the user's provided username/email and password.
 
     Returns:
         dict: A success message indicating successful authentication. The JWT token is not returned in the response body but is set in a secure cookie.
     """
-    user = authenticate_user(form_data.username, form_data.password)
+    form_data = CustomOAuth2PasswordRequestForm(
+        username=username, email=email, password=password
+    )
+    print("form_data", form_data)
+    user = authenticate_user(form_data.username, form_data.email, form_data.password)
     if not user:
         raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
@@ -154,7 +162,7 @@ async def register(response: Response, user: UserCreate):
         user_manager = UserManager(session)
 
         # Ensure unique username and email
-        existing_user = user_manager.get_user(username=user.username)
+        existing_user = user_manager.get_user_by_username(username=user.username)
         if existing_user:
             raise HTTPException(status_code=400, detail="Username already registered")
         existing_email = user_manager.get_email(email=user.email)

backend/security.py

@@ -3,6 +3,7 @@
 from fastapi.security import OAuth2PasswordBearer
 from jose import jwt, JWTError
 from passlib.context import CryptContext
+from pydantic import EmailStr
 from typing import Optional
 
 from databases.database_manager import DatabaseManager
@@ -21,12 +22,22 @@
 pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
 
 
-def authenticate_user(username: str, password: str) -> User:
+def authenticate_user(username: str, email: EmailStr, password: str) -> User:
     with DatabaseManager() as session:
         manager = UserManager(session)
-        user = manager.get_user(username)
+
+        if email:
+            print("email", email)
+            user = manager.get_user_by_email(email=email)
+            print("user", user)
+        elif username:
+            user = manager.get_user_by_username(username=username)
+        else:
+            return None
+
         if user and verify_password(password, user.hashed_password):
             return user
+
         return None
 
 
@@ -67,7 +78,7 @@ def verify_refresh_token(refresh_token: str = Cookie(None)) -> User:
         # Find the user in the database
         with DatabaseManager() as session:
             manager = UserManager(session)
-            user = manager.get_user(username)
+            user = manager.get_user_by_username(username)
 
         if user is None or user.refresh_token != refresh_token:
             raise HTTPException(
@@ -164,7 +175,7 @@ def get_current_user(request: Request) -> User:
 
     with DatabaseManager() as session:
         user_manager = UserManager(session)
-        user = user_manager.get_user(username=username)
+        user = user_manager.get_user_by_username(username=username)
 
         if user is None:
             raise HTTPException(

frontend/package.json

@@ -23,7 +23,8 @@
     "qs": "^6.11.2",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",
-    "react-router-dom": "^6.19.0"
+    "react-router-dom": "^6.19.0",
+    "validator": "^13.11.0"
   },
   "devDependencies": {
     "@types/react": "^18.2.37",

frontend/src/api/axiosInterceptor.jsx

@@ -9,7 +9,8 @@ axios.interceptors.response.use(response => response, async (error) => {
     if (error.response 
         && error.response.status === 401 
         && !originalRequest._retry 
-        && !originalRequest.url.includes('refresh-token/')) { // Check if the failed request is not for the refresh-token endpoint
+        && !originalRequest.url.includes('refresh-token/') // Check if the failed request is not for the refresh-token endpoint
+        && !originalRequest.url.includes('token/')) { // Exclude the login endpoint
         originalRequest._retry = true;
         try {
             // Attempt to refresh the token

frontend/src/pages/Login.jsx

@@ -2,26 +2,31 @@ import React, { useState } from 'react';
 import axios from 'axios';
 import qs from 'qs';
 import { Box, Button, Checkbox, Container, FormControlLabel, TextField, Typography } from '@mui/material';
+import Alert from '@mui/material/Alert';
 import LockOutlinedIcon from '@mui/icons-material/LockOutlined';
 import { useAuth } from '../contexts/AuthContext';
 import { useNavigate } from 'react-router-dom';
+import validator from 'validator';
 import { API_URL } from '../utils/constants';
 
 function LoginPage({ onLogin }) {
-  const [email, setEmail] = useState('');
+  const [usernameOrEmail, setUsernameOrEmail] = useState('');
   const [password, setPassword] = useState('');
   const [rememberMe, setRememberMe] = useState(false);
   const navigate = useNavigate();
   const { updateAuth } = useAuth();
+  const [errorMessage, setErrorMessage] = useState(''); 
 
   const handleSubmit = async (event) => {
     event.preventDefault();
 
+    // Determine if usernameOrEmail should be sent as username or email
+    const isEmail = validator.isEmail(usernameOrEmail);
+    console.log(isEmail)
+    const data = isEmail ? { email: usernameOrEmail, password } : { username: usernameOrEmail, password };
+
     try {
-        const response = await axios.post(`${API_URL}token/`, qs.stringify ({
-            username: email,
-            password
-        }), {
+        const response = await axios.post(`${API_URL}token/`, qs.stringify (data), {
             headers: {
                 'Content-Type': 'application/x-www-form-urlencoded'
             }
@@ -32,7 +37,12 @@ function LoginPage({ onLogin }) {
             navigate('/dashboards'); 
         }
     } catch (error) {
-        console.error('Login error:', error);
+        if (error.response && error.response.status === 401) {
+            // Handle 401 error here
+            setErrorMessage('Invalid credentials');
+        } else {
+            setErrorMessage(`Login error: ${error.message}`);
+        }
     }
   };
 
@@ -60,12 +70,12 @@ function LoginPage({ onLogin }) {
                   required
                   fullWidth
                   id="email"
-                  label="Email Address"
-                  name="email"
-                  autoComplete="email"
+                  label="Username or Email"
+                  name="usernameOrEmail"
+                  autoComplete="username"
                   autoFocus
-                  value={email}
-                  onChange={(e) => setEmail(e.target.value)}
+                  value={usernameOrEmail}
+                  onChange={(e) => setUsernameOrEmail(e.target.value)}
               />
               <TextField
                   margin="normal"
@@ -85,6 +95,7 @@ function LoginPage({ onLogin }) {
                   checked={rememberMe}
                   onChange={(e) => setRememberMe(e.target.checked)}
               />
+              {errorMessage && <Alert severity="error">{errorMessage}</Alert>}
               <Button
                   type="submit"
                   fullWidth

frontend/src/pages/Logout.jsx

@@ -11,6 +11,7 @@ function Logout() {
   useEffect(() => {
     // Clear the authentication cookie
     Cookies.remove('access_token');
+    Cookies.remove('refresh_token');
 
     // Update authentication state
     updateAuth(false);