Update models/api_key.py

models/api_key.py

@@ -1,345 +1,336 @@
-"""
-API密钥模型 - 处理API密钥的CRUD操作
-"""
-import json
-import uuid
-from datetime import datetime
-import os
-import pytz
-import sqlite3
-from utils.db import get_db_connection
-from config import API_KEYS_FILE, DATABASE_PATH
-
-class ApiKeyManager:
-    """管理API密钥的类"""
-    
-    @staticmethod
-    def load_keys():
-        """加载所有API密钥 (兼容旧的JSON方式)"""
-        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密钥数据 (兼容旧的JSON方式)"""
-        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():
-        """获取所有密钥"""
-        conn = get_db_connection()
-        try:
-            cursor = conn.cursor()
-            cursor.execute('SELECT * FROM api_keys')
-            rows = cursor.fetchall()
-            
-            # 转换为字典列表
-            api_keys = []
-            for row in rows:
-                key_dict = dict(row)
-                # 转换success字段从整数为布尔值
-                key_dict['success'] = bool(key_dict['success'])
-                api_keys.append(key_dict)
-            
-            return {"api_keys": api_keys}
-        except sqlite3.Error as e:
-            print(f"获取所有密钥时出错: {str(e)}")
-            # 如果数据库出错，尝试从JSON文件加载
-            return ApiKeyManager.load_keys()
-        finally:
-            conn.close()
-    
-    @staticmethod
-    def add_key(platform, name, key):
-        """添加新的API密钥"""
-        # 过滤掉key中的单引号、双引号、小括号、方括号和空格，防止存储时出错
-        if key:
-            key = key.replace("'", "").replace('"', "").replace('(', "").replace(')', "").replace('[', "").replace(']', "").replace(' ', "")
-        
-        current_time = datetime.now(pytz.timezone('Asia/Shanghai')).isoformat()
-        new_key_id = str(uuid.uuid4())
-        
-        new_key = {
-            "id": new_key_id,
-            "platform": platform,
-            "name": name,
-            "key": key,
-            "created_at": current_time,
-            "updated_at": current_time,
-            "success": False,
-            "return_message": "等待测试",
-            "states": "",
-            "balance": 0
-        }
-        
-        conn = get_db_connection()
-        try:
-            cursor = conn.cursor()
-            cursor.execute('''
-            INSERT INTO api_keys 
-            (id, platform, name, key, created_at, updated_at, success, return_message, states, balance)
-            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
-            ''', (
-                new_key_id,
-                platform,
-                name,
-                key,
-                current_time,
-                current_time,
-                0,  # success是整数: 0表示False
-                "等待测试",
-                "",
-                0
-            ))
-            conn.commit()
-            
-            # 自动调用update函数验证密钥
-            from core.update import update
-            try:
-                print(f"[自动验证] 正在验证密钥 ID: {new_key_id}")
-                update_result = update(new_key_id)
-                print(f"[自动验证] 密钥验证结果: {update_result}")
-            except Exception as e:
-                print(f"[自动验证] 验证密钥时出错: {str(e)}")
-                
-            return new_key
-        except sqlite3.Error as e:
-            print(f"添加密钥时出错: {str(e)}")
-            conn.rollback()
-            # 如果数据库出错，尝试使用JSON方式保存
-            api_keys_data = ApiKeyManager.load_keys()
-            api_keys_data["api_keys"].append(new_key)
-            ApiKeyManager.save_keys(api_keys_data)
-            return new_key
-        finally:
-            conn.close()
-    
-    @staticmethod
-    def delete_key(key_id):
-        """删除指定的API密钥"""
-        conn = get_db_connection()
-        try:
-            cursor = conn.cursor()
-            cursor.execute('DELETE FROM api_keys WHERE id = ?', (key_id,))
-            deleted = cursor.rowcount > 0
-            conn.commit()
-            return deleted
-        except sqlite3.Error as e:
-            print(f"删除密钥时出错: {str(e)}")
-            conn.rollback()
-            # 如果数据库出错，尝试从JSON文件删除
-            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
-        finally:
-            conn.close()
-    
-    @staticmethod
-    def bulk_delete(key_ids):
-        """批量删除多个API密钥"""
-        if not key_ids:
-            return 0
-        
-        conn = get_db_connection()
-        try:
-            cursor = conn.cursor()
-            # 获取当前的密钥数量
-            cursor.execute('SELECT COUNT(*) FROM api_keys')
-            original_count = cursor.fetchone()[0]
-            
-            # 使用参数化查询构建占位符
-            placeholders = ','.join(['?'] * len(key_ids))
-            cursor.execute(f'DELETE FROM api_keys WHERE id IN ({placeholders})', key_ids)
-            
-            # 获取删除后的密钥数量
-            cursor.execute('SELECT COUNT(*) FROM api_keys')
-            new_count = cursor.fetchone()[0]
-            
-            conn.commit()
-            return original_count - new_count
-        except sqlite3.Error as e:
-            print(f"批量删除密钥时出错: {str(e)}")
-            conn.rollback()
-            # 如果数据库出错，尝试从JSON文件删除
-            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
-        finally:
-            conn.close()
-    
-    @staticmethod
-    def bulk_add_keys(keys_data):
-        """批量添加多个API密钥
-        
-        Args:
-            keys_data: 包含多个密钥信息的列表，每个元素包含platform、name、key
-            
-        Returns:
-            添加的密钥列表
-        """
-        if not keys_data:
-            return []
-        
-        added_keys = []
-        now = datetime.now(pytz.timezone('Asia/Shanghai')).isoformat()
-        
-        conn = get_db_connection()
-        try:
-            cursor = conn.cursor()
-            
-            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())
-                
-                new_key = {
-                    "id": new_key_id,
-                    "platform": platform,
-                    "name": name,
-                    "key": key,
-                    "created_at": now,
-                    "updated_at": now,
-                    "success": False,
-                    "return_message": "等待测试",
-                    "states": "",
-                    "balance": 0
-                }
-                
-                cursor.execute('''
-                INSERT INTO api_keys 
-                (id, platform, name, key, created_at, updated_at, success, return_message, states, balance)
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
-                ''', (
-                    new_key_id,
-                    platform,
-                    name,
-                    key,
-                    now,
-                    now,
-                    0,  # success是整数: 0表示False
-                    "等待测试",
-                    "",
-                    0
-                ))
-                
-                added_keys.append(new_key)
-            
-            conn.commit()
-            
-            # 在批量添加完成后，启动一个单独的线程来验证所有新添加的密钥
-            from core.update import update
-            from threading import Thread
-            
-            def validate_keys():
-                for new_key in added_keys:
-                    try:
-                        key_id = new_key["id"]
-                        print(f"[自动验证] 正在验证密钥 ID: {key_id}")
-                        update_result = update(key_id)
-                        print(f"[自动验证] 密钥验证结果: {update_result}")
-                    except Exception as e:
-                        print(f"[自���验证] 验证密钥时出错: {str(e)}")
-            
-            # 启动验证线程
-            print(f"[自动验证] 启动验证线程，验证 {len(added_keys)} 个新添加的密钥")
-            validation_thread = Thread(target=validate_keys)
-            validation_thread.daemon = True  # 设置为守护线程，当主线程退出时自动结束
-            validation_thread.start()
-            
-            return added_keys
-        except sqlite3.Error as e:
-            print(f"批量添加密钥时出错: {str(e)}")
-            conn.rollback()
-            # 如果数据库出错，尝试使用JSON方式保存
-            api_keys_data = ApiKeyManager.load_keys()
-            for key in added_keys:
-                api_keys_data["api_keys"].append(key)
-            ApiKeyManager.save_keys(api_keys_data)
-            return added_keys
-        finally:
-            conn.close()
-    
-    @staticmethod
-    def update_key(key_id, name, key):
-        """更新API密钥信息"""
-        # 过滤掉key中的单引号、双引号、小括号、方括号和空格，防止存储时出错
-        if key:
-            key = key.replace("'", "").replace('"', "").replace('(', "").replace(')', "").replace('[', "").replace(']', "").replace(' ', "")
-        
-        updated_at = datetime.now(pytz.timezone('Asia/Shanghai')).isoformat()
-        
-        conn = get_db_connection()
-        try:
-            cursor = conn.cursor()
-            
-            # 更新密钥
-            cursor.execute('''
-            UPDATE api_keys 
-            SET name = ?, key = ?, updated_at = ?, success = ?, return_message = ?
-            WHERE id = ?
-            ''', (name, key, updated_at, 0, "等待测试", key_id))
-            
-            if cursor.rowcount > 0:
-                conn.commit()
-                
-                # 获取更新后的密钥信息
-                cursor.execute('SELECT * FROM api_keys WHERE id = ?', (key_id,))
-                row = cursor.fetchone()
-                
-                if row:
-                    updated_key = dict(row)
-                    # 转换success字段从整数为布尔值
-                    updated_key['success'] = bool(updated_key['success'])
-                    return updated_key
-            
-            return None
-        except sqlite3.Error as e:
-            print(f"更新密钥时出错: {str(e)}")
-            conn.rollback()
-            
-            # 如果数据库出错，尝试使用JSON方式更新
-            api_keys_data = ApiKeyManager.load_keys()
-            updated_key = None
-            for k in api_keys_data["api_keys"]:
-                if k.get("id") == key_id:
-                    k["name"] = name
-                    k["key"] = key
-                    k["updated_at"] = updated_at
-                    # 重置验证状态
-                    k["success"] = False
-                    k["return_message"] = "等待测试"
-                    updated_key = k
-                    break
-            
-            if updated_key:
-                ApiKeyManager.save_keys(api_keys_data)
-                
-            return updated_key
-        finally:
-            conn.close()
+"""
+API密钥模型 - 处理API密钥的CRUD操作
+"""
+import json
+import uuid
+from datetime import datetime
+import os
+import pytz
+import sqlite3
+from utils.db import get_db_connection
+from config import API_KEYS_FILE, DATABASE_PATH
+
+class ApiKeyManager:
+    """管理API密钥的类"""
+    
+    @staticmethod
+    def load_keys():
+        """加载所有API密钥 (兼容旧的JSON方式)"""
+        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密钥数据 (兼容旧的JSON方式)"""
+        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():
+        """获取所有密钥"""
+        conn = get_db_connection()
+        try:
+            cursor = conn.cursor()
+            cursor.execute('SELECT * FROM api_keys')
+            rows = cursor.fetchall()
+            
+            # 转换为字典列表
+            api_keys = []
+            for row in rows:
+                key_dict = dict(row)
+                # 转换success字段从整数为布尔值
+                key_dict['success'] = bool(key_dict['success'])
+                api_keys.append(key_dict)
+            
+            return {"api_keys": api_keys}
+        except sqlite3.Error as e:
+            print(f"获取所有密钥时出错: {str(e)}")
+            # 如果数据库出错，尝试从JSON文件加载
+            return ApiKeyManager.load_keys()
+        finally:
+            conn.close()
+    
+    @staticmethod
+    def add_key(platform, name, key):
+        """添加新的API密钥"""
+        # 过滤掉key中的单引号、双引号、小括号、方括号和空格，防止存储时出错
+        if key:
+            key = key.replace("'", "").replace('"', "").replace('(', "").replace(')', "").replace('[', "").replace(']', "").replace(' ', "")
+        
+        current_time = datetime.now(pytz.timezone('Asia/Shanghai')).isoformat()
+        new_key_id = str(uuid.uuid4())
+        
+        new_key = {
+            "id": new_key_id,
+            "platform": platform,
+            "name": name,
+            "key": key,
+            "created_at": current_time,
+            "updated_at": current_time,
+            "success": False,
+            "return_message": "等待测试",
+            "states": "",
+            "balance": 0
+        }
+        
+        conn = get_db_connection()
+        try:
+            cursor = conn.cursor()
+            cursor.execute('''
+            INSERT INTO api_keys 
+            (id, platform, name, key, created_at, updated_at, success, return_message, states, balance)
+            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+            ''', (
+                new_key_id,
+                platform,
+                name,
+                key,
+                current_time,
+                current_time,
+                0,  # success是整数: 0表示False
+                "等待测试",
+                "",
+                0
+            ))
+            conn.commit()
+                
+            return new_key
+        except sqlite3.Error as e:
+            print(f"添加密钥时出错: {str(e)}")
+            conn.rollback()
+            # 如果数据库出错，尝试使用JSON方式保存
+            api_keys_data = ApiKeyManager.load_keys()
+            api_keys_data["api_keys"].append(new_key)
+            ApiKeyManager.save_keys(api_keys_data)
+            return new_key
+        finally:
+            conn.close()
+    
+    @staticmethod
+    def delete_key(key_id):
+        """删除指定的API密钥"""
+        conn = get_db_connection()
+        try:
+            cursor = conn.cursor()
+            cursor.execute('DELETE FROM api_keys WHERE id = ?', (key_id,))
+            deleted = cursor.rowcount > 0
+            conn.commit()
+            return deleted
+        except sqlite3.Error as e:
+            print(f"删除密钥时出错: {str(e)}")
+            conn.rollback()
+            # 如果数据库出错，尝试从JSON文件删除
+            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
+        finally:
+            conn.close()
+    
+    @staticmethod
+    def bulk_delete(key_ids):
+        """批量删除多个API密钥"""
+        if not key_ids:
+            return 0
+        
+        conn = get_db_connection()
+        try:
+            cursor = conn.cursor()
+            # 获取当前的密钥数量
+            cursor.execute('SELECT COUNT(*) FROM api_keys')
+            original_count = cursor.fetchone()[0]
+            
+            # 使用参数化查询构建占位符
+            placeholders = ','.join(['?'] * len(key_ids))
+            cursor.execute(f'DELETE FROM api_keys WHERE id IN ({placeholders})', key_ids)
+            
+            # 获取删除后的密钥数量
+            cursor.execute('SELECT COUNT(*) FROM api_keys')
+            new_count = cursor.fetchone()[0]
+            
+            conn.commit()
+            return original_count - new_count
+        except sqlite3.Error as e:
+            print(f"批量删除密钥时出错: {str(e)}")
+            conn.rollback()
+            # 如果数据库出错，尝试从JSON文件删除
+            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
+        finally:
+            conn.close()
+    
+    @staticmethod
+    def bulk_add_keys(keys_data):
+        """批量添加多个API密钥
+        
+        Args:
+            keys_data: 包含多个密钥信息的列表，每个元素包含platform、name、key
+            
+        Returns:
+            添加的密钥列表
+        """
+        if not keys_data:
+            return []
+        
+        added_keys = []
+        now = datetime.now(pytz.timezone('Asia/Shanghai')).isoformat()
+        
+        conn = get_db_connection()
+        try:
+            cursor = conn.cursor()
+            
+            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())
+                
+                new_key = {
+                    "id": new_key_id,
+                    "platform": platform,
+                    "name": name,
+                    "key": key,
+                    "created_at": now,
+                    "updated_at": now,
+                    "success": False,
+                    "return_message": "等待测试",
+                    "states": "",
+                    "balance": 0
+                }
+                
+                cursor.execute('''
+                INSERT INTO api_keys 
+                (id, platform, name, key, created_at, updated_at, success, return_message, states, balance)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                ''', (
+                    new_key_id,
+                    platform,
+                    name,
+                    key,
+                    now,
+                    now,
+                    0,  # success是整数: 0表示False
+                    "等待测试",
+                    "",
+                    0
+                ))
+                
+                added_keys.append(new_key)
+            
+            conn.commit()
+            
+            # 在批量添加完成后，启动一个单独的线程来验证所有新添加的密钥
+            from core.update import update
+            from threading import Thread
+            
+            def validate_keys():
+                for new_key in added_keys:
+                    try:
+                        key_id = new_key["id"]
+                        print(f"[自动验证] 正在验证密钥 ID: {key_id}")
+                        update_result = update(key_id)
+                        print(f"[自动验证] 密钥验证结果: {update_result}")
+                    except Exception as e:
+                        print(f"[自动验证] 验证密钥时出错: {str(e)}")
+            
+            # 启动验证线程
+            print(f"[自动验证] 启动验证线程，验证 {len(added_keys)} 个新添加的密钥")
+            validation_thread = Thread(target=validate_keys)
+            validation_thread.daemon = True  # 设置为守护线程，当主线程退出时自动结束
+            validation_thread.start()
+            
+            return added_keys
+        except sqlite3.Error as e:
+            print(f"批量添加密钥时出错: {str(e)}")
+            conn.rollback()
+            # 如果数据库出错，尝试使用JSON方式保存
+            api_keys_data = ApiKeyManager.load_keys()
+            for key in added_keys:
+                api_keys_data["api_keys"].append(key)
+            ApiKeyManager.save_keys(api_keys_data)
+            return added_keys
+        finally:
+            conn.close()
+    
+    @staticmethod
+    def update_key(key_id, name, key):
+        """更新API密钥信息"""
+        # 过滤掉key中的单引号、双引号、小括号、方括号和空格，防止存储时出错
+        if key:
+            key = key.replace("'", "").replace('"', "").replace('(', "").replace(')', "").replace('[', "").replace(']', "").replace(' ', "")
+        
+        updated_at = datetime.now(pytz.timezone('Asia/Shanghai')).isoformat()
+        
+        conn = get_db_connection()
+        try:
+            cursor = conn.cursor()
+            
+            # 更新密钥
+            cursor.execute('''
+            UPDATE api_keys 
+            SET name = ?, key = ?, updated_at = ?, success = ?, return_message = ?
+            WHERE id = ?
+            ''', (name, key, updated_at, 0, "等待测试", key_id))
+            
+            if cursor.rowcount > 0:
+                conn.commit()
+                
+                # 获取更新后的密钥信息
+                cursor.execute('SELECT * FROM api_keys WHERE id = ?', (key_id,))
+                row = cursor.fetchone()
+                
+                if row:
+                    updated_key = dict(row)
+                    # 转换success字段从整数为布尔值
+                    updated_key['success'] = bool(updated_key['success'])
+                    return updated_key
+            
+            return None
+        except sqlite3.Error as e:
+            print(f"更新密钥时出错: {str(e)}")
+            conn.rollback()
+            
+            # 如果数据库出错，尝试使用JSON方式更新
+            api_keys_data = ApiKeyManager.load_keys()
+            updated_key = None
+            for k in api_keys_data["api_keys"]:
+                if k.get("id") == key_id:
+                    k["name"] = name
+                    k["key"] = key
+                    k["updated_at"] = updated_at
+                    # 重置验证状态
+                    k["success"] = False
+                    k["return_message"] = "等待测试"
+                    updated_key = k
+                    break
+            
+            if updated_key:
+                ApiKeyManager.save_keys(api_keys_data)
+                
+            return updated_key
+        finally:
+            conn.close()