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



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 '/login' in 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:
            return rc.response_code_view(status.HTTP_401_UNAUTHORIZED, "Token错误或失效，请重新验证")