feat: Token middleware

This commit is contained in:
fit2cloud-chenyw
2025-04-28 15:00:37 +08:00
parent 6dd8631d4d
commit b009779c0e
6 changed files with 86 additions and 4 deletions
+28
View File
@@ -0,0 +1,28 @@
from fastapi import Depends
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from common.core.config import settings
# from common.core.deps import get_current_user
from common.utils.whitelist import whiteUtils
class TokenMiddleware(BaseHTTPMiddleware):
def __init__(self, app):
super().__init__(app)
async def dispatch(self, request, call_next):
tokenkey = settings.TOKEN_KEY
if self.is_options(request) or whiteUtils.is_whitelisted(request.url.path):
return await call_next(request)
token = request.headers.get(tokenkey)
if not token or not token.startswith("Bearer "):
return JSONResponse({"error": "Unauthorized"}, status_code=401)
""" user = await get_current_user()
request.state.user = user """
return await call_next(request)
def is_options(self, request):
return request.method == "OPTIONS"
+1 -1
View File
@@ -21,7 +21,7 @@ reusable_oauth2 = XOAuth2PasswordBearer(
SessionDep = Annotated[Session, Depends(get_session)]
TokenDep = Annotated[str, Depends(reusable_oauth2)]
def get_current_user(session: SessionDep, token: TokenDep) -> sys_user:
async def get_current_user(session: SessionDep, token: TokenDep) -> sys_user:
try:
payload = jwt.decode(
token, settings.SECRET_KEY, algorithms=[security.ALGORITHM]
+54
View File
@@ -0,0 +1,54 @@
# app/utils/whitelist.py
import re
from typing import List, Pattern
from common.core.config import settings
wlist = [
"/docs",
"/login/*",
"*.json",
"*.ico"
]
class WhitelistChecker:
def __init__(self, paths: List[str] = None):
self.whitelist = paths or wlist
self._compiled_patterns: List[Pattern] = []
self._compile_patterns()
def _compile_patterns(self) -> None:
for pattern in self.whitelist:
if "*" in pattern:
regex_pattern = (
pattern.replace(".", r"\.")
.replace("*", ".*")
)
if not pattern.startswith("/"):
regex_pattern = f"^{regex_pattern}$"
else:
regex_pattern = f"^{regex_pattern}$"
try:
self._compiled_patterns.append(re.compile(regex_pattern))
except re.error:
print(f"Invalid regular expression pattern: {pattern}")
def is_whitelisted(self, path: str) -> bool:
prefix = settings.API_V1_STR
if path.startswith(prefix):
path = path[len(prefix):]
if path in self.whitelist:
return True
path = path.rstrip('/')
return any(
pattern.match(path) is not None
for pattern in self._compiled_patterns
)
def add_path(self, path: str) -> None:
if path not in self.whitelist:
self.whitelist.append(path)
if "*" in path:
self._compile_patterns()
whiteUtils = WhitelistChecker()
+2 -2
View File
@@ -3,9 +3,9 @@ from fastapi import FastAPI
from fastapi.routing import APIRoute
from starlette.middleware.cors import CORSMiddleware
from apps.api import api_router
from apps.system.middleware.auth import TokenMiddleware
from common.core.config import settings
def custom_generate_unique_id(route: APIRoute) -> str:
return f"{route.tags[0]}-{route.name}"
@@ -28,7 +28,7 @@ if settings.all_cors_origins:
allow_methods=["*"],
allow_headers=["*"],
)
app.add_middleware(TokenMiddleware)
app.include_router(api_router, prefix=settings.API_V1_STR)
if __name__ == "__main__":
+1 -1
View File
@@ -67,7 +67,7 @@ class HttpService {
// Add auth token
const token = wsCache.get('user.token')
if (token && config.headers) {
config.headers['X-SQLBOT_TOKEN'] = `Bearer ${token}`
config.headers['X-SQLBOT-TOKEN'] = `Bearer ${token}`
}
// Request logging