Linux 脚本：查询 API Key 的余额

import requests
import cmd
import readline

class KeyBalanceCLI(cmd.Cmd):
    intro = "Welcome to the Key Balance CLI. Type 'help' for a list of commands."
    prompt = "(key_balance) > "

    def do_query(self, arg):
        """
        查询单个 API Key 的余额。
        用法: query <API_KEY>
        """
        if not arg:
            print("Error: Please provide an API Key.")
            return
        balance, total_balance = get_balances(arg)
        if balance is not None and total_balance is not None:
            print_balances(arg, balance, total_balance)
        else:
            print(f"Failed to query key: {arg}")

    def do_batch_query(self, arg):
        """
        从 keys.db 文件中批量查询所有 API Key 的余额。
        用法: batch_query
        """
        try:
            with open("keys.db", "r") as file:
                keys = file.read().strip().split(",")  # 按逗号分隔读取 Key
        except FileNotFoundError:
            print("Error: keys.db file not found.")
            return
        
        total_balance = 0  # balance 总和
        total_total_balance = 0  # totalBalance 总和
        
        for key in keys:
            key = key.strip()  # 去除前后空格
            if not key:
                continue  # 跳过空 Key
            
            balance, total_balance_field = get_balances(key)  # 查询 balance 和 totalBalance
            if balance is not None and total_balance_field is not None:
                print_balances(key, balance, total_balance_field)
                
                # 累加到总和
                total_balance += balance
                total_total_balance += total_balance_field
        
        # 打印总 balance 和 totalBalance
        print(f"总 Balance: {total_balance:.2f}")
        print(f"总 TotalBalance: {total_total_balance:.2f}")

    def do_exit(self, arg):
        """
        退出程序。
        用法: exit
        """
        print("Exiting...")
        return True

def get_balances(api_key):
    """
    查询给定 API Key 的 balance 和 totalBalance。
    :param api_key: API Key
    :return: (balance, totalBalance)，如果查询失败则返回 (None, None)
    """
    url = "https://api.siliconflow.cn/v1/user/info"  # 替换为实际的 API 地址
    headers = {"Authorization": f"Bearer {api_key}"}
    
    try:
        response = requests.get(url, headers=headers)
        response.raise_for_status()  # 检查 HTTP 状态码是否为 200
        data = response.json()  # 解析 JSON 响应
        balance = float(data.get("balance", 0))  # 提取 balance 字段
        total_balance = float(data.get("totalBalance", 0))  # 提取 totalBalance 字段
        return balance, total_balance
    except requests.exceptions.RequestException as e:
        print(f"Key {api_key} 查询失败: {e}")
        return None, None

def print_balances(key, balance, total_balance):
    """
    打印 API Key 的 balance 和 totalBalance。
    :param key: API Key
    :param balance: balance 值
    :param total_balance: totalBalance 值
    """
    print(f"Key: {key}, Balance: {balance:.2f}, TotalBalance: {total_balance:.2f}")

if __name__ == "__main__":
    KeyBalanceCLI().cmdloop()