zhiyun-autoweb/main.py

# main.py

import argparse
from browser_manager import get_chrome_driver
from modules.or_address_handler import process_addresses
from modules.or_publickey_handler import generate_or_address

# 模块映射
MODULES = {
    "or-address": process_addresses,
    "or-publickey": generate_or_address,
    # 在这里可以添加更多的模块
    # "another-module": another_handler_function,
}

def main():
    """
    主函数，用于解析参数并执行相应的自动化任务。
    """
    # 1. 设置命令行参数解析
    parser = argparse.ArgumentParser(description="Selenium 自动化脚本启动器")
    parser.add_argument(
        "--module",
        type=str,
        required=True,
        choices=MODULES.keys(),
        help="要执行的业务模块名称。"
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help="使用此参数连接到已存在的 Chrome 调试实例。"
    )
    parser.add_argument(
        "--cnt",
        type=int,
        default=1,
        help="生成地址的数量。"
    )
    args = parser.parse_args()

    # 2. 根据参数选择浏览器启动方式
    use_debugger = args.debug
    driver = get_chrome_driver(use_debugger=use_debugger)
    
    if not driver:
        print("无法获取 WebDriver，程序退出。")
        return

    try:
        # 3. 根据参数调用相应的业务模块
        module_name = args.module
        handler_function = MODULES.get(module_name)
        
        if handler_function:
            print(f"--- 开始执行模块: {module_name} ---")
            handler_function(driver, args)
            print(f"--- 模块: {module_name} 执行完毕 ---")
        else:
            # 这段代码理论上不会执行，因为 argparse 的 choices 已经做了限制
            print(f"错误: 未知的模块 '{module_name}'。")
        
    except Exception as e:
        print(f"执行自动化任务时发生严重错误: {e}")
    finally:
        # 4. 任务结束后的清理工作
        print("所有任务执行完毕。")
        # 如果不是连接的调试模式，则关闭浏览器
        if not use_debugger:
            print("关闭浏览器...")
            driver.quit()
        else:
            print("任务完成，保持调试浏览器开启。")

if __name__ == "__main__":
    main()