perf: login module

Author: fit2cloud-chenyw

.gitignore

@@ -181,3 +181,5 @@ cython_debug/
 
 # PyPI configuration file
 .pypirc
+
+.DS_Store

backend/apps/system/api/login.py

@@ -1,22 +1,24 @@
-from fastapi import APIRouter, HTTPException
-from apps.system.schemas.auth import LocalLoginSchema
+from typing import Annotated
+from fastapi import APIRouter, Depends, HTTPException
+from fastapi.security import OAuth2PasswordRequestForm
 from common.core.deps import SessionDep
 from ..crud.user import authenticate
 from common.core.security import create_access_token
 from datetime import timedelta
 from common.core.config import settings
+from common.core.schemas import Token
+router = APIRouter(tags=["login"], prefix="/login")
 
-router = APIRouter(tags=["login"])
-
-@router.post("/localLogin")
+@router.post("/access-token")
 def local_login(
     session: SessionDep,
-    schema: LocalLoginSchema,
-) -> str:
-    user = authenticate(session=session, account=schema.account, password=schema.password)
+    form_data: Annotated[OAuth2PasswordRequestForm, Depends()]
+) -> Token:
+    user = authenticate(session=session, account=form_data.username, password=form_data.password)
     if not user:
         raise HTTPException(status_code=400, detail="Incorrect account or password")
     access_token_expires = timedelta(minutes=settings.ACCESS_TOKEN_EXPIRE_MINUTES)
-    return create_access_token(
-        user.id, expires_delta=access_token_expires
-    )
+    user_dict = user.to_dict()
+    return Token(access_token=create_access_token(
+        user_dict, expires_delta=access_token_expires
+    ))

backend/apps/system/api/user.py

@@ -5,7 +5,5 @@
 
 
 @router.get("/info")
-async def user_info():
-    return {
-        "name": "admin"
-    }
+async def user_info(current_user: CurrentUser):
+    return current_user.to_dict()

backend/apps/system/middleware/__init__.py

@@ -6,3 +6,10 @@ class sys_user(SQLModel, table=True):
     account: str = Field(max_length=255, unique=True)
     password: str = Field(max_length=255)
     oid: int = Field(default=1)
+
+    def to_dict(self):
+        return {
+            "id": self.id,
+            "account": self.account,
+            "oid": self.oid
+        }

backend/apps/system/middleware/token.py

@@ -2,72 +2,43 @@
 
 import jwt
 from fastapi import Depends, HTTPException, status
-from fastapi.security import OAuth2PasswordBearer
+# from fastapi.security import OAuth2PasswordBearer
 from jwt.exceptions import InvalidTokenError
-from pydantic import ValidationError
+from pydantic import BaseModel, ValidationError
 from sqlmodel import Session
-
+from common.core.schemas import TokenPayload, XOAuth2PasswordBearer
 from common.core import security
 from common.core.config import settings
 from common.core.db import get_session
 from apps.system.models.user import sys_user
-import threading
-""" reusable_oauth2 = OAuth2PasswordBearer(
+reusable_oauth2 = XOAuth2PasswordBearer(
     tokenUrl=f"{settings.API_V1_STR}/login/access-token"
 )
- """
 
-thread_local = threading.local()
+
+    
 
 SessionDep = Annotated[Session, Depends(get_session)]
-# TokenDep = Annotated[str, Depends(reusable_oauth2)]
+TokenDep = Annotated[str, Depends(reusable_oauth2)]
 
-def get_current_user(session: SessionDep, token: str) -> sys_user:
+def get_current_user(session: SessionDep, token: TokenDep) -> sys_user:
     try:
         payload = jwt.decode(
             token, settings.SECRET_KEY, algorithms=[security.ALGORITHM]
         )
-        # token_data = TokenPayload(**payload)
+        token_data = TokenPayload(**payload)
     except (InvalidTokenError, ValidationError):
         raise HTTPException(
             status_code=status.HTTP_403_FORBIDDEN,
             detail="Could not validate credentials",
         )
-    user = session.get(sys_user, payload.sub)
+    user = session.get(sys_user, token_data.id)
     if not user:
         raise HTTPException(status_code=404, detail="User not found")
     """ if not user.is_active:
         raise HTTPException(status_code=400, detail="Inactive user") """
-    thread_local.current_user = user
     return user
-def get_current_user_from_thread():
-    return getattr(thread_local, "current_user", None)
-CurrentUser = Annotated[sys_user, Depends(get_current_user_from_thread)]
-""" def get_current_user(session: SessionDep, token: TokenDep) -> User:
-    try:
-        payload = jwt.decode(
-            token, settings.SECRET_KEY, algorithms=[security.ALGORITHM]
-        )
-        token_data = TokenPayload(**payload)
-    except (InvalidTokenError, ValidationError):
-        raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
-            detail="Could not validate credentials",
-        )
-    user = session.get(User, token_data.sub)
-    if not user:
-        raise HTTPException(status_code=404, detail="User not found")
-    if not user.is_active:
-        raise HTTPException(status_code=400, detail="Inactive user")
-    return user
-
+CurrentUser = Annotated[sys_user, Depends(get_current_user)]
 
-CurrentUser = Annotated[User, Depends(get_current_user)]
 
 
-def get_current_active_superuser(current_user: CurrentUser) -> User:
-    if not current_user.is_superuser:
-        raise HTTPException(
-            status_code=403, detail="The user doesn't have enough privileges"
-        )
-    return current_user """

backend/apps/system/models/user.py

@@ -0,0 +1,31 @@
+from typing import Optional
+from fastapi import HTTPException, Request
+from fastapi.security import OAuth2PasswordBearer
+from pydantic import BaseModel
+from sqlmodel import SQLModel
+from starlette.status import HTTP_403_FORBIDDEN, HTTP_401_UNAUTHORIZED
+from common.core.config import settings
+from fastapi.security.utils import get_authorization_scheme_param
+class TokenPayload(BaseModel):
+    account: str | None = None
+    id: int | None = None
+    oid: int | None = None
+    
+class Token(SQLModel):
+    access_token: str
+    token_type: str = "bearer"
+    
+class XOAuth2PasswordBearer(OAuth2PasswordBearer):
+    async def __call__(self, request: Request) -> Optional[str]:
+        authorization = request.headers.get(settings.TOKEN_KEY)
+        scheme, param = get_authorization_scheme_param(authorization)
+        if not authorization or scheme.lower() != "bearer":
+            if self.auto_error:
+                raise HTTPException(
+                    status_code=HTTP_401_UNAUTHORIZED,
+                    detail="Not authenticated",
+                    headers={"WWW-Authenticate": "Bearer"},
+                )
+            else:
+                return None
+        return param

backend/common/core/deps.py

@@ -12,9 +12,11 @@
 ALGORITHM = "HS256"
 
 
-def create_access_token(subject: str | Any, expires_delta: timedelta) -> str:
+def create_access_token(data: dict | Any, expires_delta: timedelta) -> str:
+    to_encode = data.copy()
     expire = datetime.now(timezone.utc) + expires_delta
-    to_encode = {"exp": expire, "account": str(subject)}
+    to_encode.update({"exp": expire})
+    # to_encode = {"exp": expire, "account": str(subject)}
     encoded_jwt = jwt.encode(to_encode, settings.SECRET_KEY, algorithm=ALGORITHM)
     return encoded_jwt
 

backend/common/core/schemas.py

@@ -2,11 +2,9 @@
 from fastapi import FastAPI
 from fastapi.routing import APIRoute
 from starlette.middleware.cors import CORSMiddleware
-from starlette.middleware.base import BaseHTTPMiddleware
 from apps.api import api_router
 from common.core.config import settings
 
-from apps.system.middleware.token import verify_token
 
 def custom_generate_unique_id(route: APIRoute) -> str:
     return f"{route.tags[0]}-{route.name}"
@@ -32,7 +30,6 @@ def custom_generate_unique_id(route: APIRoute) -> str:
     )
 
 app.include_router(api_router, prefix=settings.API_V1_STR)
-# app.add_middleware(BaseHTTPMiddleware, dispatch=verify_token)
 
 if __name__ == "__main__":
     import uvicorn