models/api_key.py

"""
API密钥模型 - 处理API密钥的CRUD操作
"""
import json
import uuid
from datetime import datetime
import os
import pytz
from config import API_KEYS_FILE

class ApiKeyManager:
    """管理API密钥的类"""
    
    @staticmethod
    def load_keys():
        """加载所有API密钥"""
        if not os.path.exists(API_KEYS_FILE):
            with open(API_KEYS_FILE, 'w', encoding='utf-8') as f:
                json.dump({"api_keys": []}, f, ensure_ascii=False, indent=2)
            return {"api_keys": []}
        
        try:
            with open(API_KEYS_FILE, 'r', encoding='utf-8') as f:
                return json.load(f)
        except json.JSONDecodeError:
            return {"api_keys": []}
    
    @staticmethod
    def save_keys(data):
        """保存API密钥数据"""
        with open(API_KEYS_FILE, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
    
    @staticmethod
    def get_all_keys():
        """获取所有密钥"""
        return ApiKeyManager.load_keys()
    
    @staticmethod
    def add_key(platform, name, key):
        """添加新的API密钥"""
        api_keys_data = ApiKeyManager.load_keys()
        
        # 过滤掉key中的单引号、双引号、小括号、方括号和空格，防止存储时出错
        if key:
            key = key.replace("'", "").replace('"', "").replace('(', "").replace(')', "").replace('[', "").replace(']', "").replace(' ', "")
        
        new_key = {
            "id": str(uuid.uuid4()),
            "platform": platform,
            "name": name,
            "key": key,
            "created_at": datetime.now(pytz.timezone('Asia/Shanghai')).isoformat()
        }
        
        api_keys_data["api_keys"].append(new_key)
        ApiKeyManager.save_keys(api_keys_data)
        
        return new_key
    
    @staticmethod
    def delete_key(key_id):
        """删除指定的API密钥"""
        api_keys_data = ApiKeyManager.load_keys()
        
        original_count = len(api_keys_data["api_keys"])
        api_keys_data["api_keys"] = [k for k in api_keys_data["api_keys"] if k.get("id") != key_id]
        
        if len(api_keys_data["api_keys"]) < original_count:
            ApiKeyManager.save_keys(api_keys_data)
            return True
        return False
    
    @staticmethod
    def bulk_delete(key_ids):
        """批量删除多个API密钥"""
        if not key_ids:
            return 0
        
        api_keys_data = ApiKeyManager.load_keys()
        
        original_count = len(api_keys_data["api_keys"])
        api_keys_data["api_keys"] = [k for k in api_keys_data["api_keys"] if k.get("id") not in key_ids]
        
        deleted_count = original_count - len(api_keys_data["api_keys"])
        if deleted_count > 0:
            ApiKeyManager.save_keys(api_keys_data)
        
        return deleted_count
    
    @staticmethod
    def bulk_add_keys(keys_data):
        """批量添加多个API密钥
        
        Args:
            keys_data: 包含多个密钥信息的列表，每个元素包含platform、name、key
            
        Returns:
            添加的密钥列表
        """
        if not keys_data:
            return []
        
        api_keys_data = ApiKeyManager.load_keys()
        added_keys = []
        
        now = datetime.now(pytz.timezone('Asia/Shanghai')).isoformat()
        
        for key_info in keys_data:
            platform = key_info.get("platform")
            name = key_info.get("name")
            key = key_info.get("key")
            
            # 过滤掉key中的单引号、双引号、小括号、方括号和空格，防止存储时出错
            if key:
                key = key.replace("'", "").replace('"', "").replace('(', "").replace(')', "").replace('[', "").replace(']', "").replace(' ', "")
            
            new_key = {
                "id": str(uuid.uuid4()),
                "platform": platform,
                "name": name,
                "key": key,
                "created_at": now
            }
            
            api_keys_data["api_keys"].append(new_key)
            added_keys.append(new_key)
        
        # 一次性保存所有添加的密钥
        ApiKeyManager.save_keys(api_keys_data)
        
        return added_keys
    
    @staticmethod
    def update_key(key_id, name, key):
        """更新API密钥信息"""
        api_keys_data = ApiKeyManager.load_keys()
        
        # 过滤掉key中的单引号、双引号、小括号、方括号和空格，防止存储时出错
        if key:
            key = key.replace("'", "").replace('"', "").replace('(', "").replace(')', "").replace('[', "").replace(']', "").replace(' ', "")
            
        updated_key = None
        for k in api_keys_data["api_keys"]:
            if k.get("id") == key_id:
                k["name"] = name
                k["key"] = key
                updated_key = k
                break
        
        if updated_key:
            ApiKeyManager.save_keys(api_keys_data)
            
        return updated_key