token store

custom_auth.py

@@ -0,0 +1,78 @@
+from fastapi import Depends, HTTPException, status, Header, Query
+from typing import Optional
+from database import get_users
+from models import User, UserInDB
+from token_store import token_store
+
+
+async def get_token(
+    authorization: Optional[str] = Header(None),
+    token: Optional[str] = Query(
+        None, description="Access token (alternative to Authorization header)"
+    ),
+) -> str:
+    """
+    Extract token from Authorization header or query parameter
+    Supports both methods for better compatibility with various clients
+    """
+    # First try to get token from Authorization header
+    if authorization:
+        if authorization.startswith("Bearer "):
+            return authorization.replace("Bearer ", "")
+        else:
+            # If it doesn't have Bearer prefix, use as is
+            return authorization
+
+    # Then try to get token from query parameter
+    if token:
+        return token
+
+    # If no token is provided, raise an error
+    raise HTTPException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        detail="Authorization header missing",
+        headers={"WWW-Authenticate": "Bearer"},
+    )
+
+
+async def get_current_user_from_token(token: str = Depends(get_token)):
+    """
+    Validate token and return user if valid
+    """
+    credentials_exception = HTTPException(
+        status_code=status.HTTP_401_UNAUTHORIZED,
+        detail="Could not validate credentials",
+        headers={"WWW-Authenticate": "Bearer"},
+    )
+
+    # Validate token
+    username = token_store.validate_token(token)
+    if not username:
+        print(f"Invalid or expired token")
+        raise credentials_exception
+
+    # Get user from database
+    users = get_users()
+    if username not in users:
+        print(f"User not found: {username}")
+        raise credentials_exception
+
+    user_dict = users[username]
+    user = UserInDB(**user_dict)
+    print(f"User authenticated: {user.username}")
+
+    return user
+
+
+def create_token_for_user(username: str) -> str:
+    """
+    Create a new token for a user
+    """
+    return token_store.create_token(username)
+
+
+def remove_token(token: str) -> bool:
+    """
+    Remove a token from the store
+    """
+    return token_store.remove_token(token)

routes/auth.py

@@ -1,9 +1,13 @@
 from fastapi import APIRouter, Depends, HTTPException, status
 from fastapi.security import OAuth2PasswordRequestForm
-from datetime import timedelta
-from auth import authenticate_user, create_access_token
-from models import Token, UserCreate
-from config import ACCESS_TOKEN_EXPIRE_HOURS
+from auth import authenticate_user, get_current_user
+from custom_auth import (
+    create_token_for_user,
+    get_current_user_from_token,
+    get_token,
+    remove_token,
+)
+from models import Token, UserCreate, User
 from database import get_users, create_account
 
 router = APIRouter()
@@ -11,7 +15,7 @@ router = APIRouter()
 @router.post("/token", response_model=Token)
 async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()):
     """
-    Endpoint để lấy access token bằng username và password
+    Endpoint to get an access token using username and password
     """
     user = authenticate_user(get_users(), form_data.username, form_data.password)
     if not user:
@@ -21,19 +25,47 @@ async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends(
             headers={"WWW-Authenticate": "Bearer"},
         )
 
-    access_token_expires = timedelta(hours=ACCESS_TOKEN_EXPIRE_HOURS)
-    access_token = create_access_token(
-        data={"sub": user.username}, expires_delta=access_token_expires
-    )
+    # Create a new token for the user (this will remove any existing token)
+    access_token = create_token_for_user(user.username)
+
     return Token(access_token=access_token, token_type="bearer")
 
 
 @router.post("/register")
 async def register_user(user_data: UserCreate):
     """
-    Endpoint để đăng ký tài khoản mới
+    Endpoint to register a new account
     """
     success, message = create_account(user_data.username, user_data.password)
     if not success:
         raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=message)
     return {"message": message}
+
+
+@router.post("/logout")
+async def logout(
+    current_user: User = Depends(get_current_user_from_token),
+    token: str = Depends(get_token),
+):
+    """
+    Endpoint to logout (invalidate the current token)
+    """
+    success = remove_token(token)
+    if success:
+        return {"message": "Logout successful"}
+    else:
+        return {"message": "Token already invalid or expired"}
+
+
+@router.get("/me")
+async def get_current_user_info(
+    current_user: User = Depends(get_current_user_from_token),
+):
+    """
+    Get the current user's information
+    """
+    return {
+        "username": current_user.username,
+        "email": current_user.email,
+        "full_name": current_user.full_name,
+    }

routes/health.py

@@ -1,5 +1,5 @@
 from fastapi import APIRouter, Depends
-from simple_auth import get_current_user_from_api_key
+from custom_auth import get_current_user_from_token
 
 router = APIRouter()
 
@@ -12,7 +12,7 @@ async def health_check():
 
 
 @router.get("/auth-check")
-async def auth_check(current_user=Depends(get_current_user_from_api_key)):
+async def auth_check(current_user=Depends(get_current_user_from_token)):
     """
     Debug endpoint to verify authentication is working
     """

routes/predict.py

@@ -4,8 +4,7 @@ import shutil
 from pathlib import Path
 from fastapi import APIRouter, UploadFile, File, HTTPException, Depends, Body
 from fastapi.responses import FileResponse
-from auth import get_current_user
-from simple_auth import get_current_user_from_api_key
+from custom_auth import get_current_user_from_token
 from services.sentence_transformer_service import SentenceTransformerService, sentence_transformer_service
 from data_lib.input_name_data import InputNameData
 from data_lib.base_name_data import COL_NAME_SENTENCE
@@ -27,14 +26,14 @@ router = APIRouter()
 
 @router.post("/predict")
 async def predict(
-    current_user=Depends(get_current_user_from_api_key),
+    current_user=Depends(get_current_user_from_token),
     file: UploadFile = File(...),
     sentence_service: SentenceTransformerService = Depends(
         lambda: sentence_transformer_service
     ),
 ):
     """
-    Process an input CSV file and return standardized names (requires API Key authentication)
+    Process an input CSV file and return standardized names (requires authentication)
     """
     if not file.filename.endswith(".csv"):
         raise HTTPException(status_code=400, detail="Only CSV files are supported")
@@ -120,13 +119,13 @@ async def predict(
 @router.post("/embeddings")
 async def create_embeddings(
     request: EmbeddingRequest,
-    current_user=Depends(get_current_user_from_api_key),
+    current_user=Depends(get_current_user_from_token),
     sentence_service: SentenceTransformerService = Depends(
         lambda: sentence_transformer_service
     ),
 ):
     """
-    Create embeddings for a list of input sentences (requires API Key authentication)
+    Create embeddings for a list of input sentences (requires authentication)
     """
     try:
         start_time = time.time()
@@ -147,13 +146,13 @@ async def create_embeddings(
 @router.post("/predict-raw", response_model=PredictRawResponse)
 async def predict_raw(
     request: PredictRawRequest,
-    current_user=Depends(get_current_user_from_api_key),
+    current_user=Depends(get_current_user_from_token),
     sentence_service: SentenceTransformerService = Depends(
         lambda: sentence_transformer_service
     ),
 ):
     """
-    Process raw input records and return standardized names (requires API Key authentication)
+    Process raw input records and return standardized names (requires authentication)
     """
     try:
         # Convert input records to DataFrame

token_store.py

@@ -0,0 +1,159 @@
+import json
+import os
+import time
+import secrets
+import threading
+from datetime import datetime, timedelta
+from config import ACCESS_TOKEN_EXPIRE_HOURS
+
+
+# Singleton to store tokens
+class TokenStore:
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls):
+        with cls._lock:
+            if cls._instance is None:
+                cls._instance = super(TokenStore, cls).__new__(cls)
+                # Initialize here to make instance attributes
+                cls._instance.tokens = {}  # username -> {token, created_at}
+                cls._instance.token_to_user = {}  # token -> username
+                cls._instance.tokens_file = "data/tokens.json"
+            return cls._instance
+
+    def __init__(self):
+        # Re-initialize in __init__ to help linters recognize these attributes
+        if not hasattr(self, "tokens"):
+            self.tokens = {}
+        if not hasattr(self, "token_to_user"):
+            self.token_to_user = {}
+        if not hasattr(self, "tokens_file"):
+            self.tokens_file = "data/tokens.json"
+
+        # Load tokens when instance is created
+        if not hasattr(self, "_loaded"):
+            self._load_tokens()
+            self._loaded = True
+
+    def _load_tokens(self):
+        """Load tokens from file if it exists"""
+        os.makedirs("data", exist_ok=True)
+        if os.path.exists(self.tokens_file):
+            try:
+                with open(self.tokens_file, "r") as f:
+                    data = json.load(f)
+                    self.tokens = data.get("tokens", {})
+                    self.token_to_user = data.get("token_to_user", {})
+
+                    # Clean expired tokens on load
+                    self._clean_expired_tokens()
+            except Exception as e:
+                print(f"Error loading tokens: {e}")
+                self.tokens = {}
+                self.token_to_user = {}
+
+    def _save_tokens(self):
+        """Save tokens to file"""
+        try:
+            with open(self.tokens_file, "w") as f:
+                json.dump(
+                    {"tokens": self.tokens, "token_to_user": self.token_to_user},
+                    f,
+                    indent=4,
+                )
+        except Exception as e:
+            print(f"Error saving tokens: {e}")
+
+    def _clean_expired_tokens(self):
+        """Remove expired tokens"""
+        current_time = time.time()
+        expired_usernames = []
+        expired_tokens = []
+
+        # Find expired tokens
+        for username, token_data in self.tokens.items():
+            created_at = token_data.get("created_at", 0)
+            expiry_seconds = ACCESS_TOKEN_EXPIRE_HOURS * 3600
+
+            if current_time - created_at > expiry_seconds:
+                expired_usernames.append(username)
+                expired_tokens.append(token_data.get("token"))
+
+        # Remove expired tokens
+        for username in expired_usernames:
+            if username in self.tokens:
+                del self.tokens[username]
+
+        for token in expired_tokens:
+            if token in self.token_to_user:
+                del self.token_to_user[token]
+
+        # Save changes if any tokens were removed
+        if expired_tokens:
+            self._save_tokens()
+
+    def create_token(self, username):
+        """Create a new token for a user, removing any existing token"""
+        with self._lock:
+            # Clean expired tokens first
+            self._clean_expired_tokens()
+
+            # Remove old token if it exists
+            if username in self.tokens:
+                old_token = self.tokens[username].get("token")
+                if old_token in self.token_to_user:
+                    del self.token_to_user[old_token]
+
+            # Create new token
+            token = secrets.token_hex(32)  # 64 character random hex string
+            self.tokens[username] = {"token": token, "created_at": time.time()}
+            self.token_to_user[token] = username
+
+            # Save changes
+            self._save_tokens()
+
+            return token
+
+    def validate_token(self, token):
+        """Validate a token and return the username if valid"""
+        with self._lock:
+            # Clean expired tokens first
+            self._clean_expired_tokens()
+
+            # Check if token exists
+            if token not in self.token_to_user:
+                return None
+
+            username = self.token_to_user[token]
+
+            # Check if token is not expired
+            if username in self.tokens:
+                token_data = self.tokens[username]
+                created_at = token_data.get("created_at", 0)
+                current_time = time.time()
+                expiry_seconds = ACCESS_TOKEN_EXPIRE_HOURS * 3600
+
+                if current_time - created_at <= expiry_seconds:
+                    return username
+
+            # Token is expired or invalid
+            return None
+
+    def remove_token(self, token):
+        """Remove a token"""
+        with self._lock:
+            if token in self.token_to_user:
+                username = self.token_to_user[token]
+                del self.token_to_user[token]
+
+                if username in self.tokens:
+                    del self.tokens[username]
+
+                self._save_tokens()
+                return True
+            return False
+
+
+# Get the singleton instance
+token_store = TokenStore()