ai_api_web/backend/routes/tokens.py

from flask import Blueprint, request, jsonify
from flask_jwt_extended import jwt_required, get_jwt_identity
from models import db
from models.token import Token
from models.user import User

tokens_bp = Blueprint('tokens', __name__)

@tokens_bp.route('', methods=['GET'])
@jwt_required()
def get_tokens():
    """获取用户的令牌列表"""
    current_user_id = int(get_jwt_identity())
    search_keyword = request.args.get('keyword', '')
    search_key = request.args.get('key', '')
    
    query = Token.query.filter_by(user_id=current_user_id)
    
    if search_keyword:
        query = query.filter(Token.name.like(f'%{search_keyword}%'))
    
    if search_key:
        query = query.filter(Token.key.like(f'%{search_key}%'))
    
    tokens = query.order_by(Token.created_at.desc()).all()
    
    return jsonify({
        'tokens': [token.to_dict() for token in tokens]
    }), 200

@tokens_bp.route('', methods=['POST'])
@jwt_required()
def create_token():
    """创建新令牌"""
    current_user_id = int(get_jwt_identity())
    data = request.get_json()
    
    name = data.get('name', '新令牌')
    group = data.get('group', 'Claude Code 官方编程模型')
    remaining_quota = data.get('remaining_quota')
    total_quota = data.get('total_quota')
    available_models = data.get('available_models')
    ip_restriction = data.get('ip_restriction')
    expires_at = data.get('expires_at')
    
    token = Token(
        user_id=current_user_id,
        name=name,
        key=Token.generate_key(),
        group=group,
        remaining_quota=remaining_quota,
        total_quota=total_quota,
        available_models=available_models,
        ip_restriction=ip_restriction,
        expires_at=expires_at
    )
    
    try:
        db.session.add(token)
        db.session.commit()
        return jsonify({
            'message': '令牌创建成功',
            'token': token.to_dict(show_key=True)
        }), 201
    except Exception as e:
        db.session.rollback()
        return jsonify({'error': '创建失败: ' + str(e)}), 500

@tokens_bp.route('/<int:token_id>', methods=['GET'])
@jwt_required()
def get_token(token_id):
    """获取单个令牌详情"""
    current_user_id = int(get_jwt_identity())
    token = Token.query.filter_by(id=token_id, user_id=current_user_id).first()
    
    if not token:
        return jsonify({'error': '令牌不存在'}), 404
    
    return jsonify({'token': token.to_dict(show_key=True)}), 200

@tokens_bp.route('/<int:token_id>', methods=['PUT'])
@jwt_required()
def update_token(token_id):
    """更新令牌"""
    current_user_id = int(get_jwt_identity())
    token = Token.query.filter_by(id=token_id, user_id=current_user_id).first()
    
    if not token:
        return jsonify({'error': '令牌不存在'}), 404
    
    data = request.get_json()
    token.name = data.get('name', token.name)
    token.group = data.get('group', token.group)
    token.remaining_quota = data.get('remaining_quota', token.remaining_quota)
    token.total_quota = data.get('total_quota', token.total_quota)
    token.available_models = data.get('available_models', token.available_models)
    token.ip_restriction = data.get('ip_restriction', token.ip_restriction)
    token.expires_at = data.get('expires_at', token.expires_at)
    
    try:
        db.session.commit()
        return jsonify({
            'message': '更新成功',
            'token': token.to_dict(show_key=True)
        }), 200
    except Exception as e:
        db.session.rollback()
        return jsonify({'error': '更新失败: ' + str(e)}), 500

@tokens_bp.route('/<int:token_id>', methods=['DELETE'])
@jwt_required()
def delete_token(token_id):
    """删除令牌"""
    current_user_id = int(get_jwt_identity())
    token = Token.query.filter_by(id=token_id, user_id=current_user_id).first()
    
    if not token:
        return jsonify({'error': '令牌不存在'}), 404
    
    try:
        db.session.delete(token)
        db.session.commit()
        return jsonify({'message': '删除成功'}), 200
    except Exception as e:
        db.session.rollback()
        return jsonify({'error': '删除失败: ' + str(e)}), 500

@tokens_bp.route('/batch', methods=['POST'])
@jwt_required()
def batch_operation():
    """批量操作令牌"""
    current_user_id = int(get_jwt_identity())
    data = request.get_json()
    token_ids = data.get('token_ids', [])
    operation = data.get('operation')  # copy, delete, enable, disable
    
    tokens = Token.query.filter(
        Token.id.in_(token_ids),
        Token.user_id == current_user_id
    ).all()
    
    if not tokens:
        return jsonify({'error': '未找到有效的令牌'}), 404
    
    try:
        if operation == 'delete':
            for token in tokens:
                db.session.delete(token)
        elif operation == 'enable':
            for token in tokens:
                token.status = 'enabled'
        elif operation == 'disable':
            for token in tokens:
                token.status = 'disabled'
        
        db.session.commit()
        return jsonify({'message': '操作成功'}), 200
    except Exception as e:
        db.session.rollback()
        return jsonify({'error': '操作失败: ' + str(e)}), 500