diff --git a/backend/apps/system/middleware/__init__.py b/backend/apps/system/middleware/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/backend/apps/system/middleware/auth.py b/backend/apps/system/middleware/auth.py new file mode 100644 index 00000000..eff32618 --- /dev/null +++ b/backend/apps/system/middleware/auth.py @@ -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" + + \ No newline at end of file diff --git a/backend/common/core/deps.py b/backend/common/core/deps.py index a3e2fa17..61cfd14e 100644 --- a/backend/common/core/deps.py +++ b/backend/common/core/deps.py @@ -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] diff --git a/backend/common/utils/whitelist.py b/backend/common/utils/whitelist.py new file mode 100644 index 00000000..fa58d6a4 --- /dev/null +++ b/backend/common/utils/whitelist.py @@ -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() \ No newline at end of file diff --git a/backend/main.py b/backend/main.py index ca54a807..86e3e4ca 100644 --- a/backend/main.py +++ b/backend/main.py @@ -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__": diff --git a/frontend/src/utils/request.ts b/frontend/src/utils/request.ts index f6df087d..8b069339 100644 --- a/frontend/src/utils/request.ts +++ b/frontend/src/utils/request.ts @@ -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