key_storage.py

# 密钥存储模块

class KeyStorage:
    """
    密钥存储类，用于管理API密钥的存储
    包含有效密钥池和无效密钥池
    """
    
    def __init__(self):
        """
        初始化密钥存储
        """
        self.valid_keys = {}  # 有效密钥池，格式: {key: {"balance": float, "last_check": timestamp}}
        self.invalid_keys = {}  # 无效密钥池，格式: {key: {"reason": str, "last_check": timestamp}}
    
    def add_key(self, key, balance=None):
        """
        添加一个新的API密钥
        
        Args:
            key (str): API密钥
            balance (float, optional): 初始余额，如果为None则需要通过余额查询模块获取
        
        Returns:
            bool: 添加是否成功
        """
        if key in self.valid_keys or key in self.invalid_keys:
            return False  # 密钥已存在
        
        if balance is not None and float(balance) > 0:
            self.valid_keys[key] = {"balance": float(balance), "last_check": None}
        else:
            # 如果没有提供余额或余额为0，先放入有效池，等待余额查询模块检查
            self.valid_keys[key] = {"balance": None, "last_check": None}
        
        return True
    
    def update_key_status(self, key, balance, timestamp):
        """
        更新密钥状态
        
        Args:
            key (str): API密钥
            balance (float): 当前余额
            timestamp: 检查时间戳
        """
        # 如果余额大于0，放入有效池
        if float(balance) > 0:
            if key in self.invalid_keys:
                del self.invalid_keys[key]
            self.valid_keys[key] = {"balance": float(balance), "last_check": timestamp}
        # 否则放入无效池
        else:
            if key in self.valid_keys:
                del self.valid_keys[key]
            self.invalid_keys[key] = {"reason": "余额不足", "last_check": timestamp}
    
    def mark_key_invalid(self, key, reason, timestamp):
        """
        将密钥标记为无效
        
        Args:
            key (str): API密钥
            reason (str): 无效原因
            timestamp: 检查时间戳
        """
        if key in self.valid_keys:
            del self.valid_keys[key]
        self.invalid_keys[key] = {"reason": reason, "last_check": timestamp}
    
    def get_all_valid_keys(self):
        """
        获取所有有效密钥
        
        Returns:
            dict: 有效密钥字典
        """
        return self.valid_keys
    
    def get_all_invalid_keys(self):
        """
        获取所有无效密钥
        
        Returns:
            dict: 无效密钥字典
        """
        return self.invalid_keys
    
    def get_key_info(self, key):
        """
        获取指定密钥的信息
        
        Args:
            key (str): API密钥
        
        Returns:
            tuple: (状态, 信息字典)
        """
        if key in self.valid_keys:
            return ("valid", self.valid_keys[key])
        elif key in self.invalid_keys:
            return ("invalid", self.invalid_keys[key])
        else:
            return (None, None)
    
    def remove_key(self, key):
        """
        删除指定的密钥
        
        Args:
            key (str): 要删除的API密钥
        
        Returns:
            bool: 删除是否成功
        """
        if key in self.valid_keys:
            del self.valid_keys[key]
            return True
        elif key in self.invalid_keys:
            del self.invalid_keys[key]
            return True
        return False