LabelingSystem/backend/scripts/generate_admin_token.py

#!/usr/bin/env python3
"""
生成长期有效的管理员Token脚本

功能：
1. 查找管理员用户
2. 生成99999天有效期的Token
3. 输出Token并验证有效性

使用方式：
    cd backend
    python scripts/generate_admin_token.py

注意：需要在backend目录下运行，以确保正确加载配置
"""
import sys
import os

# 添加backend目录到路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from datetime import datetime, timedelta
import jwt
from config import settings
from database import get_db_connection


def find_admin_user():
    """查找管理员用户"""
    with get_db_connection() as conn:
        cursor = conn.cursor()
        cursor.execute("""
            SELECT id, username, email, role
            FROM users
            WHERE role = 'admin'
            LIMIT 1
        """)
        row = cursor.fetchone()
        if row:
            return {
                "id": row["id"],
                "username": row["username"],
                "email": row["email"],
                "role": row["role"]
            }
        return None


def create_long_term_token(user_data: dict, days: int = 99999) -> str:
    """
    创建长期有效的Token
    
    Args:
        user_data: 用户信息字典
        days: 有效天数，默认99999天
        
    Returns:
        str: JWT Token
    """
    expire = datetime.utcnow() + timedelta(days=days)
    payload = {
        "sub": user_data["id"],
        "username": user_data["username"],
        "email": user_data["email"],
        "role": user_data["role"],
        "exp": expire,
        "iat": datetime.utcnow(),
        "type": "access"
    }
    return jwt.encode(
        payload,
        settings.JWT_SECRET_KEY,
        algorithm=settings.JWT_ALGORITHM
    )


def verify_token(token: str) -> dict:
    """
    验证Token有效性
    
    Args:
        token: JWT Token
        
    Returns:
        dict: 解码后的payload
    """
    try:
        payload = jwt.decode(
            token,
            settings.JWT_SECRET_KEY,
            algorithms=[settings.JWT_ALGORITHM]
        )
        return payload
    except jwt.ExpiredSignatureError:
        raise Exception("Token已过期")
    except jwt.InvalidTokenError as e:
        raise Exception(f"Token无效: {str(e)}")


def main():
    print("=" * 60)
    print("管理员长期Token生成工具")
    print("=" * 60)
    print()
    
    # 查找管理员用户
    print("正在查找管理员用户...")
    admin_user = find_admin_user()
    
    if not admin_user:
        print("\n❌ 错误: 未找到管理员用户!")
        print("\n请先创建管理员用户，可以使用以下方式：")
        print("  1. 运行 python create_test_user.py 创建测试用户")
        print("  2. 或通过API注册用户后在数据库中将role改为admin")
        sys.exit(1)
    
    print(f"✓ 找到管理员用户: {admin_user['username']} ({admin_user['email']})")
    print()
    
    # 生成Token
    print("正在生成99999天有效期的Token...")
    token = create_long_term_token(admin_user, days=99999)
    print("✓ Token生成成功!")
    print()
    
    # 验证Token
    print("正在验证Token有效性...")
    try:
        payload = verify_token(token)
        expire_time = datetime.fromtimestamp(payload["exp"])
        print(f"✓ Token验证通过!")
        print(f"  - 用户ID: {payload['sub']}")
        print(f"  - 用户名: {payload['username']}")
        print(f"  - 角色: {payload['role']}")
        print(f"  - 过期时间: {expire_time.strftime('%Y-%m-%d %H:%M:%S')}")
    except Exception as e:
        print(f"❌ Token验证失败: {str(e)}")
        sys.exit(1)
    
    print()
    print("=" * 60)
    print("生成的管理员Token (请妥善保管):")
    print("=" * 60)
    print()
    print(token)
    print()
    print("=" * 60)
    print()
    print("使用方式:")
    print("  在HTTP请求头中添加:")
    print(f"  Authorization: Bearer {token[:50]}...")
    print()


if __name__ == "__main__":
    main()