crawler_task_management/public_function/auth.py

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"


async def verify_tk_token(request: Request, call_next):
    """鉴权中间件"""
    path = request.scope["path"]
    token = request.headers.get("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)