代码优化

main.py

@@ -8,7 +8,7 @@ from typing import Dict, Any, Optional, List
 from fastapi import FastAPI, HTTPException, Depends, Header
 from public_function.auth import verify_tk_token
 from public_function.public_func import read_config, create_logger
-from model.model import GoodsInfo, DataReceive, AccountCreate, AccountUpdate, DeviceResetData, CrawlerItem
+from model.model import GoodsInfo, DataReceive, AccountCreate, AccountUpdate, TokenItem, CrawlerItem
 
 app = FastAPI()
 app.middleware("http")(verify_tk_token)
@@ -184,19 +184,21 @@ async def get_goods_info(task_data: GoodsInfo, task_manager: Any = Depends(get_t
         raise HTTPException(status_code=503, detail="抓取商品数据失败，请重新尝试")
 
 
-@app.get("/device_reset", summary="设备重置")
-async def device_reset(task_data: DeviceResetData, reset_manager: Any = Depends(get_reset_manager)):
+@app.post("/add_token", summary="设备重置")
+async def add_token(task_data: TokenItem, task_manager: Any = Depends(get_task_manager)):
     """设备重置接口"""
     try:
         params = task_data.model_dump()
-        # result = await reset_manager.task_distribution(params)
-        result = True
+        result = task_manager.add_token(params.get("token"))
+        print(result)
+        # result = True
         if result:
+            print(f"新增token:{params.get("token")} 成功")
             return {"code": 200, "message": "<UNK>", "data": result}
-        raise HTTPException(status_code=404, detail="抓取商品数据失败，请重新尝试")
+        raise HTTPException(status_code=404, detail=f"新增token:{params.get("token")} 失败")
     except Exception as e:
-        print(f"设备重置失败，失败原因: {e}")
-        raise HTTPException(status_code=500, detail="获取数据失败；失败原因{}".format(e))
+        print(f"新增token:{params.get("token")} 失败，失败原因: {e}")
+        raise HTTPException(status_code=500, detail=f"新增token失败；失败原因{e}")
 
 
 if __name__ == '__main__':

model/model.py

@@ -32,10 +32,8 @@ class DataReceive(BaseModel):
     goods_info: Dict[str, Any] = Field(..., description="商品信息")
 
 
-class DeviceResetData(BaseModel):
-    device_id: str = Field(..., description="设备ID")
-    region: str = Field(..., description="地区")
-    app_name: str = Field(..., description="应用名称")
+class TokenItem(BaseModel):
+    token: str = Field(..., description="token")
 
 
 class CrawlerItem(BaseModel):

public_function/auth.py

@@ -9,17 +9,19 @@ 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):
     """鉴权中间件"""
-
+    path = request.scope["path"]
     token = request.headers.get("token")
-    if not redis_conn.check_field_exists('user_token', token):
-        return JSONResponse(status_code=401, content={"detail": "token 验证失败"})
+    if path.endswith("token"):
+        if token != TOKEN:
+            return JSONResponse(status_code=401, content={"detail": "token 验证失败"})
+    else:
+        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)

task_management/all_task_management.py

@@ -23,6 +23,20 @@ class AllTask:
             return result
         return None
 
+    def add_token(self, token) -> bool:
+        try:
+            return self.redis_conn.write_string_to_h_set("user_token", token, 1)
+        except Exception as e:
+            print(f"新增token：{token}失败,s失败原因为：{e}")
+        return False
+
+    def delete_token(self, token) -> bool:
+        try:
+            return self.redis_conn.write_string_to_h_set("user_token", token, 1)
+        except Exception as e:
+            print(f"新增token：{token}失败,s失败原因为：{e}")
+        return False
+
     async def deal_shopee_task(self, param: Dict[str, Any]):
         """
         处理Shopee任务的异步方法