zhiyinya/AIStudioBuildWS

1import os2import threading3import multiprocessing4import signal5import sys6import time7 8from browser.instance import run_browser_instance9from utils.logger import setup_logging10from utils.paths import cookies_dir, logs_dir11from utils.cookie_manager import CookieManager12from utils.common import clean_env_value, ensure_dir13 14# 全局变量15browser_processes = []16app_running = False17flask_app = None18 19 20def load_instance_configurations(logger):21    """22    使用CookieManager解析环境变量和cookies目录，为每个cookie来源创建独立的浏览器实例配置。23    """24    # 1. 读取所有实例共享的URL25    shared_url = clean_env_value(os.getenv("CAMOUFOX_INSTANCE_URL"))26    if not shared_url:27        logger.error("错误: 缺少环境变量 CAMOUFOX_INSTANCE_URL。所有实例需要一个共享的目标URL。")28        return None, None29 30    # 2. 读取全局设置31    global_settings = {32        "headless": clean_env_value(os.getenv("CAMOUFOX_HEADLESS")) or "virtual",33        "url": shared_url  # 所有实例都使用这个URL34    }35 36    proxy_value = clean_env_value(os.getenv("CAMOUFOX_PROXY"))37    if proxy_value:38        global_settings["proxy"] = proxy_value39 40    # 3. 使用CookieManager检测所有cookie来源41    cookie_manager = CookieManager(logger)42    sources = cookie_manager.detect_all_sources()43 44    # 检查是否有任何cookie来源45    if not sources:46        logger.error("错误: 未找到任何cookie来源（既没有JSON文件，也没有环境变量cookie）。")47        return None, None48 49    # 4. 为每个cookie来源创建实例配置50    instances = []51    for source in sources:52        if source.type == "file":53            instances.append({54                "cookie_file": source.identifier,55                "cookie_source": source56            })57        elif source.type == "env_var":58            # 从环境变量名中提取索引，如 "USER_COOKIE_1" -> 159            env_index = source.identifier.split("_")[-1]60            instances.append({61                "cookie_file": None,62                "env_cookie_index": int(env_index),63                "cookie_source": source64            })65 66    logger.info(f"将启动 {len(instances)} 个浏览器实例")67 68    return global_settings, instances69 70def start_browser_instances():71    """启动浏览器实例的核心逻辑"""72    global browser_processes, app_running73 74    log_dir = logs_dir()75    logger = setup_logging(str(log_dir / 'app.log'))76    logger.info("---------------------Camoufox 实例管理器开始启动---------------------")77 78    global_settings, instance_profiles = load_instance_configurations(logger)79    if not instance_profiles:80        logger.error("错误: 环境变量中未找到任何实例配置。")81        return82 83    for i, profile in enumerate(instance_profiles, 1):84        if not app_running:85            break86 87        final_config = global_settings.copy()88        final_config.update(profile)89 90        if 'url' not in final_config:91            logger.warning(f"警告: 跳过一个无效的配置项 (缺少 url): {profile}")92            continue93 94        cookie_source = final_config.get('cookie_source')95 96        if cookie_source:97            if cookie_source.type == "file":98                logger.info(99                    f"正在启动第 {i}/{len(instance_profiles)} 个浏览器实例 (file: {cookie_source.display_name})..."100                )101            elif cookie_source.type == "env_var":102                logger.info(103                    f"正在启动第 {i}/{len(instance_profiles)} 个浏览器实例 (env: {cookie_source.display_name})..."104                )105        else:106            logger.error(f"错误: 配置中缺少cookie_source对象")107            continue108 109        process = multiprocessing.Process(target=run_browser_instance, args=(final_config,))110        browser_processes.append(process)111        process.start()112 113        # 如果不是最后一个实例，等待30秒再启动下一个实例，避免并发启动导致的高CPU占用114        if i < len(instance_profiles):115            logger.info(f"等待 30 秒后启动下一个实例...")116            time.sleep(30)117 118    # 等待所有进程119    try:120        while app_running and browser_processes:121            for process in browser_processes[:]:122                if not process.is_alive():123                    browser_processes.remove(process)124                else:125                    process.join(timeout=1)126            time.sleep(1)127    except KeyboardInterrupt:128        logger.info("捕获到终止信号，正在关闭所有浏览器进程...")129        for process in browser_processes:130            process.terminate()131            process.join()132 133def run_standalone_mode():134    """独立模式"""135    global app_running136    app_running = True137 138    start_browser_instances()139 140def run_server_mode():141    """服务器模式"""142    global app_running, flask_app143 144    log_dir = logs_dir()145    server_logger = setup_logging(str(log_dir / 'app.log'), prefix="server")146 147    # 动态导入 Flask（只在需要时）148    try:149        from flask import Flask, jsonify150        flask_app = Flask(__name__)151    except ImportError:152        server_logger.error("错误: 服务器模式需要 Flask，请安装: pip install flask")153        return154 155    app_running = True156 157    # 在后台线程中启动浏览器实例158    browser_thread = threading.Thread(target=start_browser_instances, daemon=True)159    browser_thread.start()160 161    # 定义路由162    @flask_app.route('/health')163    def health_check():164        """健康检查端点"""165        running_count = sum(1 for p in browser_processes if p.is_alive())166        return jsonify({167            'status': 'healthy',168            'browser_instances': len(browser_processes),169            'running_instances': running_count,170            'message': f'Application is running with {running_count} active browser instances'171        })172 173    @flask_app.route('/')174    def index():175        """主页端点"""176        running_count = sum(1 for p in browser_processes if p.is_alive())177        return jsonify({178            'status': 'running',179            'browser_instances': len(browser_processes),180            'running_instances': running_count,181            'run_mode': 'server',182            'message': 'Camoufox Browser Automation is running in server mode'183        })184 185    # 禁用 Flask 的默认日志186    import logging187    log = logging.getLogger('werkzeug')188    log.setLevel(logging.ERROR)189 190    # 启动 Flask 服务器191    try:192        flask_app.run(host='0.0.0.0', port=7860, debug=False)193    except KeyboardInterrupt:194        server_logger.info("服务器正在关闭...")195 196def signal_handler(signum, frame):197    """统一的信号处理器"""198    global app_running199    logger = setup_logging(str(logs_dir() / 'app.log'), prefix="signal")200    logger.info(f"接收到信号 {signum}，正在关闭应用...")201    app_running = False202 203    # 关闭所有浏览器进程204    for process in browser_processes:205        if process.is_alive():206            process.terminate()207            try:208                process.join(timeout=5)209            except:210                process.kill()211 212    logger.info("所有进程已关闭")213    sys.exit(0)214 215def main():216    """主入口函数"""217    # 初始化必要的目录218    ensure_dir(logs_dir())219    ensure_dir(cookies_dir())220 221    # 注册信号处理器222    signal.signal(signal.SIGTERM, signal_handler)223    signal.signal(signal.SIGINT, signal_handler)224 225    # 检查运行模式环境变量226    hg_mode = os.getenv('HG', '').lower()227 228    if hg_mode == 'true':229        run_server_mode()230    else:231        run_standalone_mode()232 233if __name__ == "__main__":234    multiprocessing.freeze_support()235    main()236