from fastapi import status
from starlette.middleware.base import BaseHTTPMiddleware
from urllib.request import Request
from jwt import PyJWTError
from app.common import reponse_code as rc
from app.common import jwt_check as jc

green = ['/login', '/view_img', 'view_thumb']


def check_green(s: str):
    for url in green:
        if url in s:
            return True
    return False


class TokenMiddleware(BaseHTTPMiddleware):

    def __init__(self, app):
        super().__init__(app)

    async def dispatch(self, request: Request, call_next):
        """
        验证token中间件
        :param request: Request请求
        :param call_next:
        :return:
        """
        token = request.headers.get('Authorization')
        path = request.url.path
        if check_green(path):
            response = await call_next(request)
            return response
        if not token:
            return rc.response_code_view(status.HTTP_401_UNAUTHORIZED, "缺少Token，请重新验证")
        try:
            jc.check_token(token)
            return await call_next(request)
        except PyJWTError as error:
            print(error)
            return rc.response_code_view(status.HTTP_401_UNAUTHORIZED, "Token错误或失效，请重新验证")