diff --git a/main.py b/main.py index dfd8bd8..01e12b8 100644 --- a/main.py +++ b/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": "", "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__': diff --git a/model/model.py b/model/model.py index 8c7f762..777c2ba 100644 --- a/model/model.py +++ b/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): diff --git a/public_function/auth.py b/public_function/auth.py index 9a7ce94..b2b7ce8 100644 --- a/public_function/auth.py +++ b/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) diff --git a/task_management/all_task_management.py b/task_management/all_task_management.py index fdfe491..6f2ce04 100644 --- a/task_management/all_task_management.py +++ b/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任务的异步方法