Merge pull request #1 from CodeSignal/auth

Fixed logout with JWT bug and reorganized

Author: matheusgalvao1

app/routes/auth.py

@@ -1,18 +1,31 @@
-from flask import Blueprint, request, jsonify, session
-import jwt
-import datetime
+from flask import Blueprint, request, jsonify
 from config.auth_config import AuthMethod, AuthConfig
-from services.auth_service import generate_jwt_token, is_username_taken, add_user, signup_user, login_user, logout_user, blacklist_token, validate_refresh_token, refresh_tokens
+from services.auth_service import (
+    signup_user,
+    validate_refresh_token,
+    generate_jwt_token,
+    refresh_tokens,
+    login_jwt,
+    login_session,
+    logout_jwt,
+    logout_session
+)
 
-auth_bp = Blueprint("auth", __name__)
-auth_config = None
+# --- Configuration ---
+auth_bp = Blueprint("auth", __name__)  # Flask blueprint for auth routes
+auth_config = None  # Global configuration object set during initialization
 
 def init_auth_routes(config: AuthConfig):
     global auth_config
     auth_config = config
 
+# --- Authentication Routes ---
 @auth_bp.route("/signup", methods=["POST"])
 def signup():
+    """
+    Register a new user
+    Expects JSON: {"username": "user", "password": "pass"}
+    """
     if auth_config.auth_method == AuthMethod.API_KEY:
         return jsonify({"error": "Signup not available with API key authentication"}), 400
 
@@ -21,14 +34,60 @@ def signup():
 
 @auth_bp.route("/login", methods=["POST"])
 def login():
+    """
+    Authenticate user and return tokens (JWT) or create session
+    Expects JSON: {"username": "user", "password": "pass"}
+    """
     if auth_config.auth_method == AuthMethod.API_KEY:
         return jsonify({"error": "Login not available with API key authentication"}), 400
 
     data = request.get_json()
-    return login_user(data)
+    if not data or "username" not in data or "password" not in data:
+        return jsonify({"error": "Username and password are required"}), 400
+    
+    username = data["username"]
+    password = data["password"]
+    
+    if auth_config.auth_method == AuthMethod.JWT:
+        return login_jwt(username, password)
+    
+    return login_session(username, password)
+
+@auth_bp.route("/logout", methods=["POST"])
+def logout():
+    """
+    End user session or invalidate JWT tokens
+    For JWT: Requires Authorization header with Bearer token and refresh_token in JSON body
+    For Session: No additional requirements
+    """
+    if auth_config.auth_method == AuthMethod.API_KEY:
+        return jsonify({"error": "Logout not available with API key authentication"}), 400
+
+    if auth_config.auth_method == AuthMethod.JWT:
+        auth_header = request.headers.get('Authorization')
+        if not auth_header or not auth_header.startswith('Bearer '):
+            return jsonify({"error": "Access token is required in Authorization header"}), 401
+        access_token = auth_header.split(' ')[1]
+        
+        if not request.is_json:
+            return jsonify({"error": "Request must be JSON"}), 415
+        
+        data = request.get_json()
+        refresh_token = data.get("refresh_token")
+        if not refresh_token:
+            return jsonify({"error": "Refresh token is required in request body"}), 400
+        
+        return logout_jwt(access_token, refresh_token)
+    
+    return logout_session()
 
 @auth_bp.route("/refresh", methods=["POST"])
 def refresh_token():
+    """
+    Get new access token using refresh token
+    Expects JSON: {"refresh_token": "token"}
+    Returns: New access token and refresh token pair
+    """
     data = request.get_json()
     refresh_token = data.get("refresh_token")
     username = validate_refresh_token(refresh_token)
@@ -37,36 +96,10 @@ def refresh_token():
 
     access_token, new_refresh_token = generate_jwt_token(username)
 
-    # Remove old refresh token and store new one
     if refresh_token in refresh_tokens:
         del refresh_tokens[refresh_token]
 
     return jsonify({
         "access_token": access_token,
         "refresh_token": new_refresh_token
-    }), 200
-
-@auth_bp.route("/logout", methods=["POST"])
-def logout():
-    if auth_config.auth_method == AuthMethod.API_KEY:
-        return jsonify({"error": "Logout not available with API key authentication"}), 400
-    
-    elif auth_config.auth_method == AuthMethod.JWT:
-        auth_header = request.headers.get('Authorization')
-        if auth_header and auth_header.startswith('Bearer '):
-            token = auth_header.split(' ')[1]
-            blacklist_token(token)
-        
-        # Invalidate refresh token
-        data = request.get_json()
-        refresh_token = data.get("refresh_token")
-        if refresh_token in refresh_tokens:
-            del refresh_tokens[refresh_token]
-        
-        return jsonify({"message": "Logout successful"})
-    
-    elif auth_config.auth_method == AuthMethod.SESSION:
-        session.clear()
-        return jsonify({"message": "Logout successful"})
-    
-    return jsonify({"error": "Invalid authentication method"}), 500 
+    }), 200 

app/services/auth_service.py

@@ -5,98 +5,116 @@
 from models.user import User
 from flask import session, jsonify
 
-# Initialize auth_config as None
-auth_config = None
-
-# List to store User objects
-users = []
-
-# In-memory store for refresh tokens and blacklisted tokens
-refresh_tokens = {}
-blacklisted_tokens = set()
+# --- Configuration ---
+auth_config = None  # Global configuration object set during initialization
 
 def init_auth_service(config: AuthConfig):
     global auth_config
     auth_config = config
 
+# --- Storage ---
+users = []  # In-memory storage for user objects
+refresh_tokens = {}  # Maps refresh tokens to usernames
+blacklisted_tokens = set()  # Set of invalidated access tokens
+
+# --- User Management ---
+def is_username_taken(username):
+    """Check if a username is already registered"""
+    return any(user.username == username for user in users)
+
+def add_user(username, password):
+    """Add a new user if username is not taken"""
+    if not is_username_taken(username):
+        users.append(User(username, password))
+        return True
+    return False
+
+def validate_credentials(username, password):
+    """Verify username/password combination and return user if valid"""
+    user = next((user for user in users if user.username == username), None)
+    if not user or not user.check_password(password):
+        return None
+    return user
+
+# --- Token Management ---
 def generate_refresh_token(username):
+    """Create and store a new refresh token for a user"""
     refresh_token = secrets.token_hex(32)
     refresh_tokens[refresh_token] = username
     return refresh_token
 
 def validate_refresh_token(refresh_token):
+    """Check if refresh token is valid and return associated username"""
     return refresh_tokens.get(refresh_token)
 
 def blacklist_token(token):
+    """Invalidate an access token"""
     blacklisted_tokens.add(token)
 
 def generate_jwt_token(username):
+    """Generate a new JWT access token and refresh token pair"""
     access_token = jwt.encode(
         {
             "sub": username,
             "iat": datetime.datetime.utcnow(),
-            "exp": datetime.datetime.utcnow() + datetime.timedelta(minutes=15)  # Shorter expiry for access token
+            "exp": datetime.datetime.utcnow() + datetime.timedelta(minutes=15)
         },
         auth_config.jwt_secret,
         algorithm="HS256"
     )
     refresh_token = generate_refresh_token(username)
     return access_token, refresh_token
 
-# Function to check if a username already exists
-def is_username_taken(username):
-    return any(user.username == username for user in users)
-
-# Function to add a new user
-def add_user(username, password):
-    if not is_username_taken(username):
-        users.append(User(username, password))
-        return True
-    return False
-
+# --- Authentication Operations ---
 def signup_user(data):
+    """Register a new user with username and password"""
     if not data or "username" not in data or "password" not in data:
         return jsonify({"error": "Username and password are required"}), 400
 
     if is_username_taken(data["username"]):
         return jsonify({"error": "Username already exists"}), 400
 
     add_user(data["username"], data["password"])
-    
     return jsonify({"message": "Signup successful. Please log in to continue."}), 201
 
-
-def login_user(data):
-    if not data or "username" not in data or "password" not in data:
-        return jsonify({"error": "Username and password are required"}), 400
+def login_jwt(username, password):
+    """Authenticate user and return JWT tokens if valid"""
+    user = validate_credentials(username, password)
+    if not user:
+        return jsonify({"error": "Invalid username or password"}), 401
 
-    # Find user and validate credentials
-    user = next((user for user in users if user.username == data["username"]), None)
-    if not user or not user.check_password(data["password"]):
+    access_token, refresh_token = generate_jwt_token(username)
+    return jsonify({
+        "message": "Login successful",
+        "access_token": access_token,
+        "refresh_token": refresh_token
+    })
+
+def login_session(username, password):
+    """Authenticate user and create session if valid"""
+    user = validate_credentials(username, password)
+    if not user:
         return jsonify({"error": "Invalid username or password"}), 401
 
-    if auth_config.auth_method == AuthMethod.JWT:
-        access_token, refresh_token = generate_jwt_token(data["username"])
-        return jsonify({
-            "message": "Login successful",
-            "access_token": access_token,
-            "refresh_token": refresh_token
-        })
+    session["authenticated"] = True
+    session["username"] = username
+    return jsonify({"message": "Login successful"})
+
+def logout_jwt(access_token, refresh_token):
+    """Invalidate JWT access and refresh tokens"""
+    if not access_token or not refresh_token:
+        return jsonify({"error": "Both access token and refresh token are required"}), 400
 
-    elif auth_config.auth_method == AuthMethod.SESSION:
-        session["authenticated"] = True
-        session["username"] = data["username"]
-        return jsonify({"message": "Login successful"})
+    blacklist_token(access_token)
+    if refresh_token in refresh_tokens:
+        del refresh_tokens[refresh_token]
 
-    return jsonify({"error": "Invalid authentication method"}), 500
-
+    return jsonify({"message": "Logout successful"})
 
-def logout_user():
-    if auth_config.auth_method == AuthMethod.JWT:
-        return jsonify({"message": "Logout successful"})
-    
-    elif auth_config.auth_method == AuthMethod.SESSION:
-        session.clear()
-        return jsonify({"message": "Logout successful"})
+def logout_session():
+    """Clear user session if authenticated"""
+    if not session.get("authenticated"):
+        return jsonify({"error": "Not authenticated"}), 401
 
-    return jsonify({"error": "Invalid authentication method"}), 500 
+    session.clear()
+    return jsonify({"message": "Logout successful"})