mygo-kyx/codexTool
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(config, email_service, telegram_bot): Add dynamic GPTMail API key management and config reload capability

Browse Source 
- Add reload_config() function to dynamically reload config.toml and team.json without restart
- Implement GPTMail API key management with support for multiple keys (api_keys list)
- Add functions to manage GPTMail keys: get_gptmail_keys(), add_gptmail_key(), remove_gptmail_key()
- Add key rotation strategies: get_next_gptmail_key() (round-robin) and get_random_gptmail_key()
- Add gptmail_keys.json file support for dynamic key management
- Fix config parsing to handle include_team_owners and proxy settings in multiple locations
- Update email_service.py to use key rotation for GPTMail API calls
- Update telegram_bot.py to support dynamic key management
- Add install_service.sh script for service installation
- Update config.toml.example with new api_keys configuration option
- Improve configuration flexibility by supporting both old (api_key) and new (api_keys) formats
This commit is contained in:
kyx236
2026-01-17 05:52:05 +08:00
parent 64707768f8
commit b902922d22
5 changed files with 1067 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

570
telegram_bot.py
View File

@@ -7,7 +7,7 @@ from concurrent.futures import ThreadPoolExecutor
from functools import wraps
from typing import Optional
from telegram import Update, Bot
from telegram import Update, Bot, BotCommand
from telegram.ext import (
Application,
CommandHandler,
@@ -34,10 +34,22 @@ from config import (
S2A_API_BASE,
CPA_API_BASE,
CRS_API_BASE,
get_gptmail_keys,
add_gptmail_key,
remove_gptmail_key,
GPTMAIL_API_KEYS,
INCLUDE_TEAM_OWNERS,
reload_config,
S2A_CONCURRENCY,
S2A_PRIORITY,
S2A_GROUP_NAMES,
S2A_GROUP_IDS,
S2A_ADMIN_KEY,
)
from utils import load_team_tracker
from bot_notifier import BotNotifier, set_notifier, progress_finish
from s2a_service import s2a_get_dashboard_stats, format_dashboard_stats
from email_service import gptmail_service, unified_create_email
from logger import log
@@ -93,6 +105,13 @@ class ProvisionerBot:
("dashboard", self.cmd_dashboard),
("import", self.cmd_import),
("stock", self.cmd_stock),
("gptmail_keys", self.cmd_gptmail_keys),
("gptmail_add", self.cmd_gptmail_add),
("gptmail_del", self.cmd_gptmail_del),
("test_email", self.cmd_test_email),
("include_owners", self.cmd_include_owners),
("reload", self.cmd_reload),
("s2a_config", self.cmd_s2a_config),
]
for cmd, handler in handlers:
self.app.add_handler(CommandHandler(cmd, handler))
@@ -125,6 +144,10 @@ class ProvisionerBot:
# 运行 Bot
await self.app.initialize()
await self.app.start()
# 设置命令菜单提示
await self._set_commands()
await self.app.updater.start_polling(drop_pending_updates=True)
# 等待关闭信号
@@ -140,6 +163,36 @@ class ProvisionerBot:
"""请求关闭 Bot"""
self._shutdown_event.set()
async def _set_commands(self):
"""设置 Bot 命令菜单提示"""
commands = [
BotCommand("help", "查看帮助信息"),
BotCommand("list", "查看 team.json 账号列表"),
BotCommand("status", "查看任务处理状态"),
BotCommand("team", "查看指定 Team 详情"),
BotCommand("config", "查看系统配置"),
BotCommand("run", "处理指定 Team"),
BotCommand("run_all", "处理所有 Team"),
BotCommand("stop", "停止当前任务"),
BotCommand("logs", "查看最近日志"),
BotCommand("headless", "切换无头模式"),
BotCommand("dashboard", "查看 S2A 仪表盘"),
BotCommand("stock", "查看账号库存"),
BotCommand("import", "导入账号到 team.json"),
BotCommand("gptmail_keys", "查看 GPTMail API Keys"),
BotCommand("gptmail_add", "添加 GPTMail API Key"),
BotCommand("gptmail_del", "删除 GPTMail API Key"),
BotCommand("test_email", "测试邮箱创建功能"),
BotCommand("include_owners", "切换 Owner 入库开关"),
BotCommand("reload", "重载配置文件"),
BotCommand("s2a_config", "配置 S2A 服务参数"),
]
try:
await self.app.bot.set_my_commands(commands)
log.info("Bot 命令菜单已设置")
except Exception as e:
log.warning(f"设置命令菜单失败: {e}")
@admin_only
async def cmd_help(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""显示帮助信息"""
@@ -159,19 +212,28 @@ class ProvisionerBot:
<b>⚙️ 配置管理:</b>
/headless - 开启/关闭无头模式
/include_owners - 开启/关闭 Owner 入库
/reload - 重载配置文件 (无需重启)
<b>📊 S2A 专属:</b>
/dashboard - 查看 S2A 仪表盘
/stock - 查看账号库存
/s2a_config - 配置 S2A 参数
<b>📤 导入账号:</b>
/import - 导入账号到 team.json
或直接发送 JSON 文件
<b>📧 GPTMail 管理:</b>
/gptmail_keys - 查看所有 API Keys
/gptmail_add &lt;key&gt; - 添加 API Key
/gptmail_del &lt;key&gt; - 删除 API Key
/test_email - 测试邮箱创建
<b>💡 示例:</b>
<code>/list</code> - 查看所有待处理账号
<code>/run 0</code> - 处理第一个 Team
<code>/config</code> - 查看当前配置"""
<code>/gptmail_add my-api-key</code> - 添加 Key"""
await update.message.reply_text(help_text, parse_mode="HTML")
@admin_only
@@ -300,6 +362,9 @@ class ProvisionerBot:
# 无头模式状态
headless_status = "✅ 已开启" if BROWSER_HEADLESS else "❌ 未开启"
# Owner 入库状态
include_owners_status = "✅ 已开启" if INCLUDE_TEAM_OWNERS else "❌ 未开启"
lines = [
"<b>⚙️ 系统配置</b>",
"",
@@ -309,6 +374,7 @@ class ProvisionerBot:
"<b>🔐 授权服务</b>",
f" 模式: {AUTH_PROVIDER.upper()}",
f" 地址: {auth_url}",
f" Owner 入库: {include_owners_status}",
"",
"<b>🌐 浏览器</b>",
f" 无头模式: {headless_status}",
@@ -321,7 +387,8 @@ class ProvisionerBot:
f" 状态: {proxy_info}",
"",
"<b>💡 提示:</b>",
"使用 /headless 开启/关闭无头模式",
"/headless - 切换无头模式",
"/include_owners - 切换 Owner 入库",
]
await update.message.reply_text("\n".join(lines), parse_mode="HTML")
@@ -354,7 +421,7 @@ class ProvisionerBot:
await update.message.reply_text(
f"<b>🌐 无头模式</b>\n\n"
f"状态: {status}\n\n"
f"⚠️ 需要重启 Bot 生效",
f"💡 使用 /reload 立即生效",
parse_mode="HTML"
)
@@ -366,6 +433,246 @@ class ProvisionerBot:
except Exception as e:
await update.message.reply_text(f"❌ 修改配置失败: {e}")
@admin_only
async def cmd_include_owners(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""切换 Owner 入库开关"""
import tomli_w
try:
# 读取当前配置
with open(CONFIG_FILE, "rb") as f:
import tomllib
config = tomllib.load(f)
# 获取当前状态 (顶层配置)
current = config.get("include_team_owners", False)
new_value = not current
# 更新配置 (写到顶层)
config["include_team_owners"] = new_value
# 写回文件
with open(CONFIG_FILE, "wb") as f:
tomli_w.dump(config, f)
status = "✅ 已开启" if new_value else "❌ 已关闭"
desc = "运行任务时会将 Team Owner 账号也入库到授权服务" if new_value else "运行任务时不会入库 Team Owner 账号"
await update.message.reply_text(
f"<b>👤 Owner 入库开关</b>\n\n"
f"状态: {status}\n"
f"{desc}\n\n"
f"💡 使用 /reload 立即生效",
parse_mode="HTML"
)
except ImportError:
await update.message.reply_text(
"❌ 缺少 tomli_w 依赖\n"
"请运行: uv add tomli_w"
)
except Exception as e:
await update.message.reply_text(f"❌ 修改配置失败: {e}")
@admin_only
async def cmd_reload(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""重载配置文件"""
# 检查是否有任务正在运行
if self.current_task and not self.current_task.done():
await update.message.reply_text(
f"⚠️ 有任务正在运行: {self.current_team}\n"
"请等待任务完成或使用 /stop 停止后再重载配置"
)
return
await update.message.reply_text("⏳ 正在重载配置...")
try:
# 调用重载函数
result = reload_config()
if result["success"]:
# 重新导入更新后的配置变量
from config import (
EMAIL_PROVIDER as new_email_provider,
AUTH_PROVIDER as new_auth_provider,
INCLUDE_TEAM_OWNERS as new_include_owners,
BROWSER_HEADLESS as new_headless,
ACCOUNTS_PER_TEAM as new_accounts_per_team,
)
lines = [
"<b>✅ 配置重载成功</b>",
"",
f"📁 team.json: {result['teams_count']} 个账号",
f"📄 config.toml: {'已更新' if result['config_updated'] else '未变化'}",
"",
"<b>当前配置:</b>",
f" 邮箱服务: {new_email_provider}",
f" 授权服务: {new_auth_provider}",
f" Owner 入库: {'✅' if new_include_owners else '❌'}",
f" 无头模式: {'✅' if new_headless else '❌'}",
f" 每 Team 账号: {new_accounts_per_team}",
]
await update.message.reply_text("\n".join(lines), parse_mode="HTML")
else:
await update.message.reply_text(
f"<b>❌ 配置重载失败</b>\n\n{result['message']}",
parse_mode="HTML"
)
except Exception as e:
await update.message.reply_text(f"❌ 重载配置失败: {e}")
@admin_only
async def cmd_s2a_config(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""配置 S2A 服务参数"""
import tomli_w
# 无参数时显示当前配置
if not context.args:
# 脱敏显示 API Key
key_display = "未配置"
if S2A_ADMIN_KEY:
if len(S2A_ADMIN_KEY) > 10:
key_display = f"{S2A_ADMIN_KEY[:4]}...{S2A_ADMIN_KEY[-4:]}"
else:
key_display = S2A_ADMIN_KEY[:4] + "..."
groups_display = ", ".join(S2A_GROUP_NAMES) if S2A_GROUP_NAMES else "默认分组"
group_ids_display = ", ".join(S2A_GROUP_IDS) if S2A_GROUP_IDS else "无"
lines = [
"<b>📊 S2A 服务配置</b>",
"",
f"<b>API 地址:</b> {S2A_API_BASE or '未配置'}",
f"<b>Admin Key:</b> <code>{key_display}</code>",
f"<b>并发数:</b> {S2A_CONCURRENCY}",
f"<b>优先级:</b> {S2A_PRIORITY}",
f"<b>分组名称:</b> {groups_display}",
f"<b>分组 ID:</b> {group_ids_display}",
"",
"<b>💡 修改配置:</b>",
"<code>/s2a_config concurrency 10</code>",
"<code>/s2a_config priority 50</code>",
"<code>/s2a_config groups 分组1,分组2</code>",
"<code>/s2a_config api_base https://...</code>",
"<code>/s2a_config admin_key sk-xxx</code>",
]
await update.message.reply_text("\n".join(lines), parse_mode="HTML")
return
# 解析参数
param = context.args[0].lower()
value = " ".join(context.args[1:]) if len(context.args) > 1 else None
if not value:
await update.message.reply_text(
f"❌ 缺少值\n用法: /s2a_config {param} <值>"
)
return
try:
# 读取当前配置
with open(CONFIG_FILE, "rb") as f:
import tomllib
config = tomllib.load(f)
# 确保 s2a section 存在
if "s2a" not in config:
config["s2a"] = {}
# 根据参数类型处理
updated_key = None
updated_value = None
if param == "concurrency":
try:
new_val = int(value)
if new_val < 1 or new_val > 100:
await update.message.reply_text("❌ 并发数范围: 1-100")
return
config["s2a"]["concurrency"] = new_val
updated_key = "并发数"
updated_value = str(new_val)
except ValueError:
await update.message.reply_text("❌ 并发数必须是数字")
return
elif param == "priority":
try:
new_val = int(value)
if new_val < 0 or new_val > 100:
await update.message.reply_text("❌ 优先级范围: 0-100")
return
config["s2a"]["priority"] = new_val
updated_key = "优先级"
updated_value = str(new_val)
except ValueError:
await update.message.reply_text("❌ 优先级必须是数字")
return
elif param in ("groups", "group_names"):
# 支持逗号分隔的分组名称
groups = [g.strip() for g in value.split(",") if g.strip()]
config["s2a"]["group_names"] = groups
updated_key = "分组名称"
updated_value = ", ".join(groups) if groups else "默认分组"
elif param == "group_ids":
# 支持逗号分隔的分组 ID
ids = [i.strip() for i in value.split(",") if i.strip()]
config["s2a"]["group_ids"] = ids
updated_key = "分组 ID"
updated_value = ", ".join(ids) if ids else "无"
elif param == "api_base":
config["s2a"]["api_base"] = value
updated_key = "API 地址"
updated_value = value
elif param == "admin_key":
config["s2a"]["admin_key"] = value
updated_key = "Admin Key"
# 脱敏显示
if len(value) > 10:
updated_value = f"{value[:4]}...{value[-4:]}"
else:
updated_value = value[:4] + "..."
else:
await update.message.reply_text(
f"❌ 未知参数: {param}\n\n"
"可用参数:\n"
"• concurrency - 并发数\n"
"• priority - 优先级\n"
"• groups - 分组名称\n"
"• group_ids - 分组 ID\n"
"• api_base - API 地址\n"
"• admin_key - Admin Key"
)
return
# 写回文件
with open(CONFIG_FILE, "wb") as f:
tomli_w.dump(config, f)
await update.message.reply_text(
f"<b>✅ S2A 配置已更新</b>\n\n"
f"{updated_key}: <code>{updated_value}</code>\n\n"
f"💡 使用 /reload 立即生效",
parse_mode="HTML"
)
except ImportError:
await update.message.reply_text(
"❌ 缺少 tomli_w 依赖\n"
"请运行: uv add tomli_w"
)
except Exception as e:
await update.message.reply_text(f"❌ 修改配置失败: {e}")
@admin_only
async def cmd_run(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""启动处理指定 Team"""
@@ -452,24 +759,56 @@ class ProvisionerBot:
@admin_only
async def cmd_stop(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""停止当前任务"""
"""强制停止当前任务"""
if not self.current_task or self.current_task.done():
await update.message.reply_text("📭 当前没有运行中的任务")
return
# 注意: 由于任务在线程池中运行，无法直接取消
# 这里只能发送信号
await update.message.reply_text(
f"🛑 正在停止: {self.current_team}\n"
"注意: 当前账号处理完成后才会停止"
)
task_name = self.current_team or "未知任务"
await update.message.reply_text(f"🛑 正在强制停止: {task_name}...")
# 设置全局停止标志
try:
import run
run._shutdown_requested = True
except Exception:
pass
# 1. 设置全局停止标志
try:
import run
run._shutdown_requested = True
except Exception:
pass
# 2. 取消 asyncio 任务
if self.current_task and not self.current_task.done():
self.current_task.cancel()
# 3. 强制关闭浏览器进程
try:
from browser_automation import cleanup_chrome_processes
cleanup_chrome_processes()
except Exception as e:
log.warning(f"清理浏览器进程失败: {e}")
# 4. 重置状态
self.current_team = None
# 5. 重置停止标志 (以便下次任务可以正常运行)
try:
import run
run._shutdown_requested = False
except Exception:
pass
# 清理进度跟踪
progress_finish()
await update.message.reply_text(
f"<b>✅ 任务已强制停止</b>\n\n"
f"已停止: {task_name}\n"
f"已清理浏览器进程\n\n"
f"使用 /status 查看状态",
parse_mode="HTML"
)
except Exception as e:
await update.message.reply_text(f"❌ 停止任务时出错: {e}")
@admin_only
async def cmd_logs(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
@@ -711,10 +1050,12 @@ class ProvisionerBot:
await update.message.reply_text("❌ 未找到有效账号 (需要 account/email 和 token 字段)")
return
# 读取现有 team.json
# 读取现有 team.json (不存在则自动创建)
team_json_path = Path(TEAM_JSON_FILE)
existing_accounts = []
if team_json_path.exists():
is_new_file = not team_json_path.exists()
if not is_new_file:
try:
with open(team_json_path, "r", encoding="utf-8") as f:
existing_accounts = json.load(f)
@@ -741,16 +1082,23 @@ class ProvisionerBot:
existing_emails.add(email)
added += 1
# 保存到 team.json
# 保存到 team.json (自动创建文件)
try:
# 确保父目录存在
team_json_path.parent.mkdir(parents=True, exist_ok=True)
with open(team_json_path, "w", encoding="utf-8") as f:
json.dump(existing_accounts, f, ensure_ascii=False, indent=2)
file_status = "📄 已创建 team.json" if is_new_file else "📄 已更新 team.json"
await update.message.reply_text(
f"<b>✅ 导入完成</b>\n\n"
f"{file_status}\n"
f"新增: {added}\n"
f"跳过 (重复): {skipped}\n"
f"team.json 总数: {len(existing_accounts)}\n\n"
f"💡 使用 /reload 刷新配置\n"
f"使用 /run_all 或 /run &lt;n&gt; 开始处理",
parse_mode="HTML"
)
@@ -758,6 +1106,190 @@ class ProvisionerBot:
except Exception as e:
await update.message.reply_text(f"❌ 保存到 team.json 失败: {e}")
# ==================== GPTMail Key 管理命令 ====================
@admin_only
async def cmd_gptmail_keys(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""查看所有 GPTMail API Keys"""
keys = get_gptmail_keys()
config_keys = set(GPTMAIL_API_KEYS)
if not keys:
await update.message.reply_text(
"<b>📧 GPTMail API Keys</b>\n\n"
"📭 暂无配置 Key\n\n"
"使用 /gptmail_add &lt;key&gt; 添加",
parse_mode="HTML"
)
return
lines = [f"<b>📧 GPTMail API Keys (共 {len(keys)} 个)</b>\n"]
for i, key in enumerate(keys):
# 脱敏显示
if len(key) > 10:
masked = f"{key[:4]}...{key[-4:]}"
else:
masked = key[:4] + "..." if len(key) > 4 else key
# 标记来源
source = "📁 配置" if key in config_keys else "🔧 动态"
lines.append(f"{i+1}. <code>{masked}</code> {source}")
lines.append(f"\n<b>💡 管理:</b>")
lines.append(f"/gptmail_add &lt;key&gt; - 添加 Key")
lines.append(f"/gptmail_del &lt;key&gt; - 删除动态 Key")
lines.append(f"/test_email - 测试邮箱创建")
await update.message.reply_text("\n".join(lines), parse_mode="HTML")
@admin_only
async def cmd_gptmail_add(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""添加 GPTMail API Key (支持批量导入)"""
if not context.args:
await update.message.reply_text(
"<b>📧 添加 GPTMail API Key</b>\n\n"
"<b>单个添加:</b>\n"
"<code>/gptmail_add gpt-xxx</code>\n\n"
"<b>批量添加 (空格分隔):</b>\n"
"<code>/gptmail_add key1 key2 key3</code>\n\n"
"<b>批量添加 (换行分隔):</b>\n"
"<code>/gptmail_add key1\nkey2\nkey3</code>",
parse_mode="HTML"
)
return
# 合并所有参数，支持空格和换行分隔
raw_input = " ".join(context.args)
# 按空格和换行分割
keys = []
for part in raw_input.replace("\n", " ").split():
key = part.strip()
if key:
keys.append(key)
if not keys:
await update.message.reply_text("❌ Key 不能为空")
return
# 获取现有 keys
existing_keys = set(get_gptmail_keys())
# 统计结果
added = []
skipped = []
invalid = []
await update.message.reply_text(f"⏳ 正在验证 {len(keys)} 个 Key...")
for key in keys:
# 检查是否已存在
if key in existing_keys:
skipped.append(key)
continue
# 测试 Key 是否有效
success, message = gptmail_service.test_api_key(key)
if not success:
invalid.append(key)
continue
# 添加 Key
if add_gptmail_key(key):
added.append(key)
existing_keys.add(key)
# 生成结果报告
lines = ["<b>📧 GPTMail Key 导入结果</b>\n"]
if added:
lines.append(f"<b>✅ 成功添加:</b> {len(added)}")
for k in added[:5]: # 最多显示5个
masked = f"{k[:4]}...{k[-4:]}" if len(k) > 10 else k
lines.append(f" • <code>{masked}</code>")
if len(added) > 5:
lines.append(f" • ... 等 {len(added)} 个")
if skipped:
lines.append(f"\n<b>⏭️ 已跳过 (已存在):</b> {len(skipped)}")
if invalid:
lines.append(f"\n<b>❌ 无效 Key:</b> {len(invalid)}")
for k in invalid[:3]: # 最多显示3个
masked = f"{k[:4]}...{k[-4:]}" if len(k) > 10 else k
lines.append(f" • <code>{masked}</code>")
lines.append(f"\n<b>当前 Key 总数:</b> {len(get_gptmail_keys())}")
await update.message.reply_text("\n".join(lines), parse_mode="HTML")
@admin_only
async def cmd_gptmail_del(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""删除 GPTMail API Key"""
if not context.args:
await update.message.reply_text(
"<b>📧 删除 GPTMail API Key</b>\n\n"
"用法: /gptmail_del &lt;key&gt;\n\n"
"注意: 只能删除动态添加的 Key，配置文件中的 Key 请直接修改 config.toml",
parse_mode="HTML"
)
return
key = context.args[0].strip()
# 检查是否是配置文件中的 Key
if key in GPTMAIL_API_KEYS:
await update.message.reply_text(
"⚠️ 该 Key 在配置文件中，无法通过 Bot 删除\n"
"请直接修改 config.toml"
)
return
# 删除 Key
if remove_gptmail_key(key):
await update.message.reply_text(
f"<b>✅ Key 已删除</b>\n\n"
f"当前 Key 总数: {len(get_gptmail_keys())}",
parse_mode="HTML"
)
else:
await update.message.reply_text("❌ Key 不存在或删除失败")
@admin_only
async def cmd_test_email(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
"""测试邮箱创建功能"""
if EMAIL_PROVIDER != "gptmail":
await update.message.reply_text(
f"⚠️ 当前邮箱提供商: {EMAIL_PROVIDER}\n"
f"测试功能仅支持 GPTMail 模式"
)
return
await update.message.reply_text("⏳ 正在测试邮箱创建...")
try:
# 测试创建邮箱
email, password = unified_create_email()
if email:
await update.message.reply_text(
f"<b>✅ 邮箱创建成功</b>\n\n"
f"邮箱: <code>{email}</code>\n"
f"密码: <code>{password}</code>\n\n"
f"当前 Key 数量: {len(get_gptmail_keys())}",
parse_mode="HTML"
)
else:
await update.message.reply_text(
"<b>❌ 邮箱创建失败</b>\n\n"
"请检查 GPTMail API Key 配置",
parse_mode="HTML"
)
except Exception as e:
await update.message.reply_text(f"❌ 测试失败: {e}")
async def main():
"""主函数"""