jude-cehua/auth_middleware.py

import jwt
import time
from functools import wraps
from flask import request, jsonify, redirect, url_for, session
from http import HTTPStatus
from config import config

def is_api_request() -> bool:
    """判断是否为API请求"""
    return (request.path.startswith('/api/') or 
            request.headers.get('X-Requested-With') == 'XMLHttpRequest' or
            'application/json' in (request.headers.get('Accept') or ''))

class AuthMiddleware:
    """认证中间件"""
    
    def __init__(self, app=None):
        self.app = app
        if app is not None:
            self.init_app(app)
    
    def init_app(self, app):
        """初始化应用"""
        self.app = app
        app.config['JWT_SECRET_KEY'] = config.secret_key
        app.config['JWT_ACCESS_TOKEN_EXPIRES'] = 7 * 24 * 60 * 60  # 7天
    
    def generate_token(self, user_id: str, user_info: dict) -> str:
        """生成JWT令牌"""
        payload = {
            'user_id': user_id,
            'user_info': user_info,
            'exp': time.time() + self.app.config['JWT_ACCESS_TOKEN_EXPIRES'],
            'iat': time.time()
        }
        token = jwt.encode(payload, self.app.config['JWT_SECRET_KEY'], algorithm='HS256')
        return token.decode('utf-8') if isinstance(token, bytes) else token
    
    def verify_token(self, token: str) -> dict:
        """验证JWT令牌"""
        import logging
        logger = logging.getLogger(__name__)
        
        try:
            payload = jwt.decode(token, self.app.config['JWT_SECRET_KEY'], algorithms=['HS256'])
            return payload
        except jwt.ExpiredSignatureError:
            logger.warning('JWT token已过期')
            raise ValueError('令牌已过期')
        except jwt.InvalidTokenError as e:
            logger.warning(f'JWT token无效: {str(e)}')
            raise ValueError('无效的令牌')
    
    def login_required(self, f):
        """登录要求装饰器"""
        @wraps(f)
        def decorated_function(*args, **kwargs):
            token = self._get_token()
            
            if not token:
                return self._handle_no_token()
            
            try:
                payload = self.verify_token(token)
                request.user_id = payload['user_id']
                request.user_info = payload['user_info']
                return f(*args, **kwargs)
            except ValueError as e:
                return self._handle_invalid_token(e)
        
        return decorated_function
    
    def _get_token(self) -> str:
        """获取token"""
        # 从请求头获取
        if 'Authorization' in request.headers:
            auth_header = request.headers['Authorization']
            if auth_header.startswith('Bearer '):
                return auth_header.split(' ')[1]
        
        # 从session获取
        return session.get('auth_token')
    
    def _handle_no_token(self):
        """处理无token情况"""
        if is_api_request():
            return jsonify({
                'success': False,
                'error': '需要登录',
                'code': 'AUTH_REQUIRED'
            }), HTTPStatus.UNAUTHORIZED
        return redirect('/login')
    
    def _handle_invalid_token(self, error):
        """处理无效token情况"""
        if is_api_request():
            error_str = str(error)
            error_code = 'INVALID_TOKEN'
            
            if '令牌已过期' in error_str or 'expired' in error_str.lower():
                error_code = 'TOKEN_EXPIRED'
            elif '无效的令牌' in error_str or 'invalid' in error_str.lower():
                error_code = 'TOKEN_INVALID'
            
            return jsonify({
                'success': False,
                'error': str(error),
                'code': error_code,
                'can_refresh': 'TOKEN_EXPIRED' in error_code
            }), HTTPStatus.UNAUTHORIZED

        # 根据错误类型重定向到不同页面
        error_str = str(error)
        if '令牌已过期' in error_str or 'expired' in error_str.lower():
            return redirect('/login_error?error=token_expired&message=登录已过期，请重新登录&code=TOKEN_EXPIRED')
        elif '无效的令牌' in error_str or 'invalid' in error_str.lower():
            return redirect('/login_error?error=token_invalid&message=登录凭证无效，请重新登录&code=TOKEN_INVALID')
        else:
            return redirect('/login')

# 创建全局实例
auth_middleware = AuthMiddleware()

def init_auth(app):
    """初始化认证中间件"""
    auth_middleware.init_app(app)
    return auth_middleware