代码优化

public_function/auth.py

@@ -1,14 +1,25 @@
+import os
+from pathlib import Path
 from fastapi import Request
 from fastapi.responses import JSONResponse
+from public_function.public_func import read_config
+from public_function.redis_task_manager import RedisTaskManager
+
+config_path = os.path.join(Path(__file__).resolve().parent, 'config.yaml')
+config = read_config(config_path)
+redis_conn = RedisTaskManager(config)
+
+
+# TOKEN = "opB4ztbdw45xFoJbXti20520bsEq3UDKKAtiDWHnGjjhP6v0KNFjqBM7bfzto6GLdUPviYnVdCgdCJYqe42nPoy6mvW59F3TPQZu"
 
-TOKEN = "opB4ztbdw45xFoJbXti20520bsEq3UDKKAtiDWHnGjjhP6v0KNFjqBM7bfzto6GLdUPviYnVdCgdCJYqe42nPoy6mvW59F3TPQZu"
 
 
 async def verify_tk_token(request: Request, call_next):
     """鉴权中间件"""
 
     token = request.headers.get("token")
-    if token != TOKEN:
+    if not redis_conn.check_field_exists('user_token', token):
         return JSONResponse(status_code=401, content={"detail": "token 验证失败"})
-
+    # if token != TOKEN:
+    #     return JSONResponse(status_code=401, content={"detail": "token 验证失败"})
     return await call_next(request)

public_function/redis_task_manager.py

@@ -73,6 +73,34 @@ class RedisTaskManager:
             print(f"读取Redis数据时发生错误: {e}")
             return None
 
+    def check_field_exists(self, key: str, field: str) -> bool:
+        """检查字段是否存在"""
+        try:
+            exists = self.redis_client.hexists(key, field)
+            return exists
+        except Exception as e:
+            print(f"检查字段存在失败: {e}")
+            return False
+
+    def delete_field(self, key: str, field: str) -> bool:
+        """删除HSET中的字段"""
+        try:
+            result = self.redis_client.hdel(key, field)
+            return result > 0
+        except Exception as e:
+            print(f"删除字段失败: {e}")
+            return False
+
+    def write_string_to_h_set(self, key: str, field: str, value: str) -> bool:
+        """将字符串写入Redis HSET"""
+        # key 默认为 user_token, field 默认为 token 值,value 为 123
+        try:
+            result = self.redis_client.hset(key, field, value)
+            return result > 0
+        except Exception as e:
+            print(f"写入HSET失败: {e}")
+            return False
+
     def add_task_to_set(self, task_data: Dict[str, Any]) -> bool:
         """添加任务到Redis集合"""
         try:
@@ -173,3 +201,16 @@ class RedisTaskManager:
         except Exception as e:
             print(f"清空任务失败: {e}")
             return False
+
+
+if __name__ == '__main__':
+    import os
+    from pathlib import Path
+    from public_function.public_func import read_config
+
+    config_path = os.path.join(Path(__file__).resolve().parent, 'config.yaml')
+    config = read_config(config_path)
+    redis_conn = RedisTaskManager(config)
+    token = "opB4ztbdw45xFoJbXti20520bsEq3UDKKAtiDWHnGjjhP6v0KNFjqBM7bfzto6GLdUPviYnVdCgdCJYqe42nPoy6mvW59F3TPQZu"
+    # redis_conn.write_string_to_h_set("user_token", token, 1)
+    print(redis_conn.check_field_exists("user_token", token))