From e58936bf054e3d973e6883e2f9952abc791d63ed Mon Sep 17 00:00:00 2001
From: JOJO <1498581755@qq.com>
Date: Mon, 5 Jan 2026 14:13:51 +0800
Subject: [PATCH] chore: read balance creds from .env by default

---
 modules/balance_client.py | 48 ++++++++++++++++++++++++++++++++-------
 1 file changed, 40 insertions(+), 8 deletions(-)

diff --git a/modules/balance_client.py b/modules/balance_client.py
index dd0c56d..35c8d37 100644
--- a/modules/balance_client.py
+++ b/modules/balance_client.py
@@ -1,7 +1,7 @@
 """
 Admin balance fetchers for Kimi (Moonshot), DeepSeek, and Qwen (Aliyun BSS).
 
-Credentials (read from environment):
+Credentials (优先读取仓库根目录 .env，其次环境变量):
   - MOONSHOT_API_KEY            : Bearer token for Kimi
   - DEEPSEEK_API_KEY            : Bearer token for DeepSeek
   - ALIYUN_ACCESS_KEY_ID        : AccessKey ID for Aliyun (Qwen billing)
@@ -18,11 +18,44 @@ import json
 import os
 import time
 import uuid
+from pathlib import Path
 from typing import Any, Dict, Tuple
-from urllib import parse, request, error
+from urllib import parse, request
+
+# -------- .env 读取助手 --------
+_DOTENV_CACHE: Dict[str, str] | None = None
 
 
-USD_CNY_RATE = float(os.environ.get("USD_CNY_RATE", "7.1"))
+def _read_dotenv() -> Dict[str, str]:
+    global _DOTENV_CACHE
+    if _DOTENV_CACHE is not None:
+        return _DOTENV_CACHE
+    path = Path(__file__).resolve().parents[1] / ".env"
+    result: Dict[str, str] = {}
+    if path.exists():
+        try:
+            for raw_line in path.read_text(encoding="utf-8").splitlines():
+                line = raw_line.strip()
+                if not line or line.startswith("#") or "=" not in line:
+                    continue
+                key, value = line.split("=", 1)
+                key = key.strip()
+                value = value.strip().strip('"').strip("'")
+                if key:
+                    result[key] = value
+        except Exception:
+            pass
+    _DOTENV_CACHE = result
+    return result
+
+
+def _env(key: str, default: str | None = None) -> str | None:
+    if key in os.environ:
+        return os.environ[key]
+    return _read_dotenv().get(key, default)
+
+
+USD_CNY_RATE = float(_env("USD_CNY_RATE", "7.1") or "7.1")
 
 
 def _http_get(url: str, headers: Dict[str, str] | None = None, timeout: int = 8) -> Tuple[Dict[str, Any] | None, str | None]:
@@ -38,7 +71,7 @@ def _http_get(url: str, headers: Dict[str, str] | None = None, timeout: int = 8)
 
 # -------- Kimi (Moonshot) --------
 def fetch_kimi_balance() -> Dict[str, Any]:
-    api_key = os.environ.get("MOONSHOT_API_KEY")
+    api_key = _env("MOONSHOT_API_KEY")
     if not api_key:
         return {"success": False, "error": "MOONSHOT_API_KEY 未设置"}
 
@@ -68,7 +101,7 @@ def fetch_kimi_balance() -> Dict[str, Any]:
 
 # -------- DeepSeek --------
 def fetch_deepseek_balance() -> Dict[str, Any]:
-    api_key = os.environ.get("DEEPSEEK_API_KEY")
+    api_key = _env("DEEPSEEK_API_KEY")
     if not api_key:
         return {"success": False, "error": "DEEPSEEK_API_KEY 未设置"}
 
@@ -110,8 +143,8 @@ def _sign(params: Dict[str, Any], secret: str) -> str:
 
 
 def fetch_qwen_balance() -> Dict[str, Any]:
-    ak = os.environ.get("ALIYUN_ACCESS_KEY_ID")
-    sk = os.environ.get("ALIYUN_ACCESS_KEY_SECRET")
+    ak = _env("ALIYUN_ACCESS_KEY_ID")
+    sk = _env("ALIYUN_ACCESS_KEY_SECRET")
     if not ak or not sk:
         return {"success": False, "error": "缺少 ALIYUN_ACCESS_KEY_ID / ALIYUN_ACCESS_KEY_SECRET"}
 
@@ -156,4 +189,3 @@ def fetch_all_balances() -> Dict[str, Any]:
         "deepseek": fetch_deepseek_balance(),
         "qwen": fetch_qwen_balance(),
     }
-