561 lines
19 KiB
Python
561 lines
19 KiB
Python
import asyncio
|
||
import sys
|
||
import os
|
||
from datetime import datetime
|
||
from typing import Dict, Any, List, Callable, Awaitable
|
||
|
||
from loguru import logger
|
||
|
||
# 导入功能模块
|
||
import upload_accounts
|
||
import reset_extracted
|
||
import reupload_all
|
||
try:
|
||
import retry_failed
|
||
has_retry_failed = True
|
||
except ImportError:
|
||
has_retry_failed = False
|
||
|
||
try:
|
||
import reset_retrying
|
||
has_reset_retrying = True
|
||
except ImportError:
|
||
has_reset_retrying = False
|
||
|
||
try:
|
||
import delete_db
|
||
has_delete_db = True
|
||
except ImportError:
|
||
has_delete_db = False
|
||
|
||
from core.config import Config
|
||
from core.logger import setup_logger
|
||
from core.database import DatabaseManager
|
||
|
||
|
||
class MenuSystem:
|
||
def __init__(self):
|
||
self.config = Config.from_yaml()
|
||
self.logger = setup_logger(self.config)
|
||
self.db_manager = DatabaseManager(self.config)
|
||
|
||
# 定义菜单项
|
||
self.menu_items = [
|
||
{
|
||
"key": "1",
|
||
"title": "重新上传所有账号",
|
||
"description": "重置所有账号的提取状态并重新上传到API服务器（推荐）",
|
||
"handler": self.reupload_all_menu
|
||
},
|
||
{
|
||
"key": "2",
|
||
"title": "上传账号到API",
|
||
"description": "将注册成功的账号上传到API服务器",
|
||
"handler": self.upload_accounts_menu
|
||
},
|
||
{
|
||
"key": "3",
|
||
"title": "重置账号提取状态",
|
||
"description": "将账号的extracted字段重置为0，允许重新上传",
|
||
"handler": self.reset_extracted_menu
|
||
}
|
||
]
|
||
|
||
# 添加可选功能
|
||
if has_retry_failed:
|
||
self.menu_items.append({
|
||
"key": "4",
|
||
"title": "重试失败账号",
|
||
"description": "重新尝试注册失败的账号",
|
||
"handler": self.retry_failed_menu
|
||
})
|
||
|
||
if has_reset_retrying:
|
||
self.menu_items.append({
|
||
"key": "5",
|
||
"title": "重置正在重试的账号",
|
||
"description": "将状态为retrying的账号重置为failed",
|
||
"handler": self.reset_retrying_menu
|
||
})
|
||
|
||
if has_delete_db:
|
||
self.menu_items.append({
|
||
"key": "6",
|
||
"title": "删除数据库记录",
|
||
"description": "根据条件删除数据库中的记录",
|
||
"handler": self.delete_db_menu
|
||
})
|
||
|
||
# 添加退出选项
|
||
self.menu_items.append({
|
||
"key": "q",
|
||
"title": "退出程序",
|
||
"description": "退出菜单系统",
|
||
"handler": self.exit_menu
|
||
})
|
||
|
||
async def initialize(self):
|
||
"""初始化数据库连接"""
|
||
await self.db_manager.initialize()
|
||
|
||
async def cleanup(self):
|
||
"""清理资源"""
|
||
await self.db_manager.cleanup()
|
||
|
||
async def show_menu(self):
|
||
"""显示主菜单"""
|
||
os.system('cls' if os.name == 'nt' else 'clear')
|
||
print("\n===== Cursor 账号管理系统 =====\n")
|
||
|
||
# 显示菜单项
|
||
for item in self.menu_items:
|
||
print(f"{item['key']}. {item['title']} - {item['description']}")
|
||
|
||
print("\n请选择功能 (输入对应数字或q退出):", end=" ")
|
||
choice = input().strip().lower()
|
||
|
||
# 查找选择的菜单项
|
||
selected_item = next((item for item in self.menu_items if item["key"] == choice), None)
|
||
|
||
if selected_item:
|
||
await selected_item["handler"]()
|
||
else:
|
||
print("\n无效的选择，请重试!")
|
||
input("按Enter键继续...")
|
||
await self.show_menu()
|
||
|
||
async def upload_accounts_menu(self):
|
||
"""上传账号菜单"""
|
||
os.system('cls' if os.name == 'nt' else 'clear')
|
||
print("\n===== 上传账号到API =====\n")
|
||
|
||
# 统计账号状态
|
||
resetter = reset_extracted.ExtractedResetter()
|
||
await resetter.initialize()
|
||
stats = await resetter.count_success_accounts()
|
||
await resetter.cleanup()
|
||
|
||
print(f"当前状态: 总成功账号: {stats['total']}, 已提取: {stats['extracted']}, 未提取: {stats['not_extracted']}")
|
||
|
||
if stats['not_extracted'] == 0:
|
||
print("\n没有待上传的账号。是否要重置所有账号的提取状态？(y/n)", end=" ")
|
||
reset_choice = input().strip().lower()
|
||
if reset_choice == 'y':
|
||
await self.reset_all_extracted()
|
||
print("\n已重置所有账号的提取状态，现在可以上传。")
|
||
else:
|
||
print("\n未重置状态，无法上传账号。")
|
||
input("\n按Enter键返回主菜单...")
|
||
return
|
||
|
||
print("\n选项:")
|
||
print("1. 预览待上传账号 (dry-run)")
|
||
print("2. 上传所有账号")
|
||
print("3. 上传指定批次数量")
|
||
print("4. 修改批次大小 (默认: 100)")
|
||
print("5. 禁用代理")
|
||
print("b. 返回主菜单")
|
||
|
||
choice = input("\n请选择操作: ").strip().lower()
|
||
|
||
args = []
|
||
use_proxy = True
|
||
batch_size = 100
|
||
max_batches = 0
|
||
|
||
if choice == "1":
|
||
args = ["--dry-run"]
|
||
elif choice == "2":
|
||
args = [] # 使用默认设置上传所有
|
||
elif choice == "3":
|
||
try:
|
||
batches = int(input("请输入要处理的批次数量: ").strip())
|
||
if batches > 0:
|
||
args = [f"--batches", f"{batches}"]
|
||
max_batches = batches
|
||
else:
|
||
print("批次数量必须大于0")
|
||
input("按Enter键继续...")
|
||
await self.upload_accounts_menu()
|
||
return
|
||
except ValueError:
|
||
print("无效的输入，必须是整数")
|
||
input("按Enter键继续...")
|
||
await self.upload_accounts_menu()
|
||
return
|
||
elif choice == "4":
|
||
try:
|
||
size = int(input("请输入批次大小: ").strip())
|
||
if size > 0:
|
||
batch_size = size
|
||
args = [f"--batch-size", f"{size}"]
|
||
else:
|
||
print("批次大小必须大于0")
|
||
input("按Enter键继续...")
|
||
await self.upload_accounts_menu()
|
||
return
|
||
except ValueError:
|
||
print("无效的输入，必须是整数")
|
||
input("按Enter键继续...")
|
||
await self.upload_accounts_menu()
|
||
return
|
||
elif choice == "5":
|
||
use_proxy = False
|
||
args = ["--no-proxy"]
|
||
elif choice == "b":
|
||
await self.show_menu()
|
||
return
|
||
else:
|
||
print("无效的选择，请重试!")
|
||
input("按Enter键继续...")
|
||
await self.upload_accounts_menu()
|
||
return
|
||
|
||
# 处理输入的批次大小和批次数量组合
|
||
if choice == "3" and choice == "4":
|
||
args = [f"--batches", f"{max_batches}", f"--batch-size", f"{batch_size}"]
|
||
|
||
# 添加代理设置
|
||
if not use_proxy and "--no-proxy" not in args:
|
||
args.append("--no-proxy")
|
||
|
||
print(f"\n准备执行命令: python upload_accounts.py {' '.join(args)}")
|
||
print("按Enter键开始执行，或按Ctrl+C取消...")
|
||
input()
|
||
|
||
# 存储原始sys.argv并替换
|
||
original_argv = sys.argv.copy()
|
||
sys.argv = [sys.argv[0]] + args
|
||
|
||
try:
|
||
# 执行上传账号功能
|
||
await upload_accounts.main()
|
||
except Exception as e:
|
||
self.logger.error(f"执行上传账号出错: {e}")
|
||
finally:
|
||
# 恢复sys.argv
|
||
sys.argv = original_argv
|
||
|
||
input("\n操作完成，按Enter键返回主菜单...")
|
||
await self.show_menu()
|
||
|
||
async def reset_extracted_menu(self):
|
||
"""重置提取状态菜单"""
|
||
os.system('cls' if os.name == 'nt' else 'clear')
|
||
print("\n===== 重置账号提取状态 =====\n")
|
||
|
||
# 统计账号状态
|
||
resetter = reset_extracted.ExtractedResetter()
|
||
await resetter.initialize()
|
||
stats = await resetter.count_success_accounts()
|
||
|
||
print(f"当前状态: 总成功账号: {stats['total']}, 已提取: {stats['extracted']}, 未提取: {stats['not_extracted']}")
|
||
|
||
print("\n选项:")
|
||
print("1. 预览账号状态 (dry-run)")
|
||
print("2. 重置所有账号的提取状态")
|
||
print("3. 重置指定模式的账号提取状态")
|
||
print("b. 返回主菜单")
|
||
|
||
choice = input("\n请选择操作: ").strip().lower()
|
||
|
||
args = []
|
||
|
||
if choice == "1":
|
||
args = ["--dry-run"]
|
||
elif choice == "2":
|
||
print("\n确定要重置所有账号的提取状态吗？这将允许已提取的账号再次上传。(y/n)", end=" ")
|
||
confirm = input().strip().lower()
|
||
if confirm != 'y':
|
||
print("操作已取消")
|
||
input("按Enter键继续...")
|
||
await self.reset_extracted_menu()
|
||
return
|
||
elif choice == "3":
|
||
pattern = input("请输入邮箱匹配模式 (例如: outlook.com): ").strip()
|
||
if pattern:
|
||
args = ["--pattern", pattern]
|
||
else:
|
||
print("模式不能为空")
|
||
input("按Enter键继续...")
|
||
await self.reset_extracted_menu()
|
||
return
|
||
elif choice == "b":
|
||
await self.show_menu()
|
||
return
|
||
else:
|
||
print("无效的选择，请重试!")
|
||
input("按Enter键继续...")
|
||
await self.reset_extracted_menu()
|
||
return
|
||
|
||
# 执行重置
|
||
print(f"\n准备执行命令: python reset_extracted.py {' '.join(args)}")
|
||
print("按Enter键开始执行，或按Ctrl+C取消...")
|
||
input()
|
||
|
||
# 存储原始sys.argv并替换
|
||
original_argv = sys.argv.copy()
|
||
sys.argv = [sys.argv[0]] + args
|
||
|
||
try:
|
||
# 执行重置操作
|
||
await reset_extracted.main()
|
||
except Exception as e:
|
||
self.logger.error(f"执行重置提取状态出错: {e}")
|
||
finally:
|
||
# 恢复sys.argv
|
||
sys.argv = original_argv
|
||
await resetter.cleanup()
|
||
|
||
input("\n操作完成，按Enter键返回主菜单...")
|
||
await self.show_menu()
|
||
|
||
async def reupload_all_menu(self):
|
||
"""重新上传所有账号菜单"""
|
||
os.system('cls' if os.name == 'nt' else 'clear')
|
||
print("\n===== 重新上传所有账号 =====\n")
|
||
|
||
print("此操作将:")
|
||
print("1. 重置所有成功账号的提取状态")
|
||
print("2. 重新上传所有账号到API服务器")
|
||
|
||
print("\n警告: 此操作可能需要较长时间，请确保网络连接稳定。")
|
||
print("是否继续? (y/n)", end=" ")
|
||
|
||
choice = input().strip().lower()
|
||
if choice != 'y':
|
||
print("\n操作已取消")
|
||
input("\n按Enter键返回主菜单...")
|
||
await self.show_menu()
|
||
return
|
||
|
||
try:
|
||
# 执行重新上传
|
||
await reupload_all.main()
|
||
except Exception as e:
|
||
self.logger.error(f"重新上传过程中出错: {e}")
|
||
finally:
|
||
input("\n操作完成，按Enter键返回主菜单...")
|
||
await self.show_menu()
|
||
|
||
async def retry_failed_menu(self):
|
||
"""重试失败账号菜单"""
|
||
if not has_retry_failed:
|
||
print("\n功能不可用: retry_failed.py 未找到")
|
||
input("按Enter键返回主菜单...")
|
||
await self.show_menu()
|
||
return
|
||
|
||
os.system('cls' if os.name == 'nt' else 'clear')
|
||
print("\n===== 重试失败账号 =====\n")
|
||
|
||
print("此功能将重新尝试注册失败的账号。")
|
||
# 这里可以添加更多的统计信息显示
|
||
|
||
print("\n选项:")
|
||
print("1. 执行账号重试")
|
||
print("b. 返回主菜单")
|
||
|
||
choice = input("\n请选择操作: ").strip().lower()
|
||
|
||
if choice == "1":
|
||
# 执行重试失败账号功能
|
||
print("\n准备重试失败账号...")
|
||
print("按Enter键开始执行，或按Ctrl+C取消...")
|
||
input()
|
||
|
||
# 存储原始sys.argv
|
||
original_argv = sys.argv.copy()
|
||
sys.argv = [sys.argv[0]]
|
||
|
||
try:
|
||
# 执行重试功能
|
||
await retry_failed.main()
|
||
except Exception as e:
|
||
self.logger.error(f"执行重试失败账号出错: {e}")
|
||
finally:
|
||
# 恢复sys.argv
|
||
sys.argv = original_argv
|
||
elif choice == "b":
|
||
await self.show_menu()
|
||
return
|
||
else:
|
||
print("无效的选择，请重试!")
|
||
input("按Enter键继续...")
|
||
await self.retry_failed_menu()
|
||
return
|
||
|
||
input("\n操作完成，按Enter键返回主菜单...")
|
||
await self.show_menu()
|
||
|
||
async def reset_retrying_menu(self):
|
||
"""重置正在重试的账号菜单"""
|
||
if not has_reset_retrying:
|
||
print("\n功能不可用: reset_retrying.py 未找到")
|
||
input("按Enter键返回主菜单...")
|
||
await self.show_menu()
|
||
return
|
||
|
||
os.system('cls' if os.name == 'nt' else 'clear')
|
||
print("\n===== 重置正在重试的账号 =====\n")
|
||
|
||
print("此功能将状态为retrying的账号重置为failed。")
|
||
# 这里可以添加更多的统计信息显示
|
||
|
||
print("\n选项:")
|
||
print("1. 执行重置retrying状态")
|
||
print("b. 返回主菜单")
|
||
|
||
choice = input("\n请选择操作: ").strip().lower()
|
||
|
||
if choice == "1":
|
||
# 执行重置retrying状态功能
|
||
print("\n准备重置正在重试的账号...")
|
||
print("按Enter键开始执行，或按Ctrl+C取消...")
|
||
input()
|
||
|
||
# 存储原始sys.argv
|
||
original_argv = sys.argv.copy()
|
||
sys.argv = [sys.argv[0]]
|
||
|
||
try:
|
||
# 执行重置功能
|
||
await reset_retrying.main()
|
||
except Exception as e:
|
||
self.logger.error(f"执行重置retrying状态出错: {e}")
|
||
finally:
|
||
# 恢复sys.argv
|
||
sys.argv = original_argv
|
||
elif choice == "b":
|
||
await self.show_menu()
|
||
return
|
||
else:
|
||
print("无效的选择，请重试!")
|
||
input("按Enter键继续...")
|
||
await self.reset_retrying_menu()
|
||
return
|
||
|
||
input("\n操作完成，按Enter键返回主菜单...")
|
||
await self.show_menu()
|
||
|
||
async def delete_db_menu(self):
|
||
"""删除数据库记录菜单"""
|
||
if not has_delete_db:
|
||
print("\n功能不可用: delete_db.py 未找到")
|
||
input("按Enter键返回主菜单...")
|
||
await self.show_menu()
|
||
return
|
||
|
||
os.system('cls' if os.name == 'nt' else 'clear')
|
||
print("\n===== 删除数据库记录 =====\n")
|
||
|
||
print("此功能将根据条件删除数据库中的记录。")
|
||
print("警告: 删除操作不可恢复，请谨慎操作！")
|
||
|
||
print("\n选项:")
|
||
print("1. 交互式删除（按状态）")
|
||
print("2. 删除所有记录")
|
||
print("b. 返回主菜单")
|
||
|
||
choice = input("\n请选择操作: ").strip().lower()
|
||
|
||
if choice == "b":
|
||
await self.show_menu()
|
||
return
|
||
elif choice == "2":
|
||
print("\n警告: 此操作将删除数据库中的所有记录！")
|
||
print("确定要继续吗？(y/n)", end=" ")
|
||
confirm = input().strip().lower()
|
||
if confirm != 'y':
|
||
print("操作已取消")
|
||
input("按Enter键继续...")
|
||
await self.delete_db_menu()
|
||
return
|
||
|
||
# 执行全部删除
|
||
args = ["--all", "--confirm"]
|
||
print(f"\n准备执行命令: python delete_db.py {' '.join(args)}")
|
||
print("按Enter键开始执行，或按Ctrl+C取消...")
|
||
input()
|
||
|
||
# 存储原始sys.argv
|
||
original_argv = sys.argv.copy()
|
||
sys.argv = [sys.argv[0]] + args
|
||
|
||
try:
|
||
# 执行删除功能
|
||
await delete_db.main()
|
||
except Exception as e:
|
||
self.logger.error(f"执行删除数据库记录出错: {e}")
|
||
finally:
|
||
# 恢复sys.argv
|
||
sys.argv = original_argv
|
||
|
||
input("\n操作完成，按Enter键返回主菜单...")
|
||
await self.show_menu()
|
||
return
|
||
elif choice == "1":
|
||
# 执行交互式删除
|
||
args = ["--interactive"]
|
||
print(f"\n准备执行命令: python delete_db.py {' '.join(args)}")
|
||
print("按Enter键开始执行，或按Ctrl+C取消...")
|
||
input()
|
||
|
||
# 存储原始sys.argv
|
||
original_argv = sys.argv.copy()
|
||
sys.argv = [sys.argv[0]] + args
|
||
|
||
try:
|
||
# 执行删除功能
|
||
await delete_db.main()
|
||
except Exception as e:
|
||
self.logger.error(f"执行删除数据库记录出错: {e}")
|
||
finally:
|
||
# 恢复sys.argv
|
||
sys.argv = original_argv
|
||
|
||
input("\n操作完成，按Enter键返回主菜单...")
|
||
await self.show_menu()
|
||
return
|
||
else:
|
||
print("无效的选择，请重试!")
|
||
input("按Enter键继续...")
|
||
await self.delete_db_menu()
|
||
return
|
||
|
||
async def reset_all_extracted(self):
|
||
"""重置所有账号的提取状态"""
|
||
resetter = reset_extracted.ExtractedResetter()
|
||
await resetter.initialize()
|
||
try:
|
||
await resetter.reset_extracted()
|
||
finally:
|
||
await resetter.cleanup()
|
||
|
||
async def exit_menu(self):
|
||
"""退出菜单"""
|
||
print("\n感谢使用，再见！")
|
||
sys.exit(0)
|
||
|
||
|
||
async def main():
|
||
menu = MenuSystem()
|
||
try:
|
||
await menu.initialize()
|
||
await menu.show_menu()
|
||
except KeyboardInterrupt:
|
||
print("\n程序被用户中断")
|
||
except Exception as e:
|
||
logger.error(f"程序执行出错: {e}")
|
||
finally:
|
||
await menu.cleanup()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
try:
|
||
asyncio.run(main())
|
||
except KeyboardInterrupt:
|
||
print("\n程序被用户中断")
|
||
except Exception as e:
|
||
print(f"程序崩溃: {e}") |