core/update.py

"""
API密钥更新模块 - 提供API密钥的验证和更新功能
"""
import json
import os
import sqlite3
import pytz
from datetime import datetime
from core.api_manager import get_api_manager
from utils.db import get_db_connection
from config import API_KEYS_FILE

def update(key_id):
    """
    更新指定ID的API密钥
    
    Args:
        key_id (str): 要更新的API密钥ID
        
    Returns:
        dict: 包含更新结果的字典，成功时返回更新后的密钥信息，失败时返回错误信息
    """
    # 从SQLite数据库中获取密钥信息
    conn = get_db_connection()
    try:
        cursor = conn.cursor()
        cursor.execute('SELECT * FROM api_keys WHERE id = ?', (key_id,))
        row = cursor.fetchone()
        
        if row is None:
            # 数据库中找不到，尝试从JSON文件加载
            # 这是为了支持旧版本的兼容
            if os.path.exists(API_KEYS_FILE):
                try:
                    with open(API_KEYS_FILE, "r", encoding="utf-8") as f:
                        data = json.load(f)
                    
                    for key in data.get("api_keys", []):
                        if key.get("id") == key_id:
                            key_data = key
                            break
                    else:
                        return {"success": False, "message": f"未找到ID为 {key_id} 的API密钥"}
                except Exception as e:
                    return {"success": False, "message": f"读取API密钥文件失败: {str(e)}"}
            else:
                return {"success": False, "message": f"未找到ID为 {key_id} 的API密钥"}
        else:
            key_data = dict(row)
        
        # 获取平台和密钥
        platform = key_data.get("platform")
        api_key = key_data.get("key")
        
        # 获取API管理器
        api_manager = get_api_manager()
        
        # 调用API管理器验证密钥
        try:
            result = api_manager.execute(platform, "validate_api_key", api_key)
        except Exception as e:
            return {"success": False, "message": f"验证API密钥时出错: {str(e)}"}
        
        # 当前时间 - 使用亚洲/上海时区 (UTC+8)
        current_time = datetime.now(pytz.timezone('Asia/Shanghai')).isoformat()
        
        # 将布尔值转换为整数
        success_int = 1 if result.get("success", False) else 0
        
        # 更新密钥信息到SQLite数据库
        try:
            cursor.execute('''
            UPDATE api_keys
            SET states = ?, balance = ?, success = ?, return_message = ?, updated_at = ?
            WHERE id = ?
            ''', (
                result.get("states", ""),
                result.get("balance", 0),
                success_int,
                result.get("return_message", ""),
                current_time,
                key_id
            ))
            
            # 如果数据库中没有此记录（可能是旧的JSON格式数据），则插入新记录
            if cursor.rowcount == 0 and 'id' in key_data:
                cursor.execute('''
                INSERT OR REPLACE INTO api_keys
                (id, platform, name, key, created_at, updated_at, success, return_message, states, balance)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                ''', (
                    key_data.get("id"),
                    key_data.get("platform"),
                    key_data.get("name"),
                    key_data.get("key"),
                    key_data.get("created_at", current_time),
                    current_time,
                    success_int,
                    result.get("return_message", ""),
                    result.get("states", ""),
                    result.get("balance", 0)
                ))
            
            conn.commit()
            
            # 获取更新后的完整记录
            cursor.execute('SELECT * FROM api_keys WHERE id = ?', (key_id,))
            updated_row = cursor.fetchone()
            
            if updated_row:
                updated_data = dict(updated_row)
                # 将布尔值转换为布尔类型
                updated_data['success'] = bool(updated_data['success'])
                
                return {
                    "success": True, 
                    "message": "API密钥更新成功",
                    "data": updated_data
                }
            else:
                # 如果以某种方式在更新过程中密钥被删除
                return {"success": False, "message": f"更新期间ID为 {key_id} 的API密钥已被删除"}
            
        except sqlite3.Error as e:
            conn.rollback()
            return {"success": False, "message": f"更新数据库中的API密钥失败: {str(e)}"}
        
    except sqlite3.Error as e:
        return {"success": False, "message": f"从数据库获取API密钥时出错: {str(e)}"}
    finally:
        if conn:
            conn.close()

if __name__ == "__main__":
    # 可以在这里添加命令行参数解析的代码，用于直接从命令行调用
    import sys
    if len(sys.argv) > 1:
        key_id = sys.argv[1]
        result = update(key_id)
        print(json.dumps(result, indent=2, ensure_ascii=False))