mirror of
https://gh-proxy.org/https://github.com/klinsmaya/SQLBot.git
synced 2026-11-08 01:06:59 +08:00
feat: Token middleware
This commit is contained in:
@@ -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"
|
||||
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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
@@ -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__":
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user