feat: Token middleware

Author: fit2cloud-chenyw

backend/apps/system/middleware/__init__.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"
+    
+    

backend/apps/system/middleware/auth.py

@@ -21,7 +21,7 @@
 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]

backend/common/core/deps.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()

backend/common/utils/whitelist.py

@@ -3,9 +3,9 @@
 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 @@ def custom_generate_unique_id(route: APIRoute) -> str:
         allow_methods=["*"],
         allow_headers=["*"],
     )
-
+app.add_middleware(TokenMiddleware)
 app.include_router(api_router, prefix=settings.API_V1_STR)
 
 if __name__ == "__main__":

backend/main.py

@@ -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