# --- browser_utils/initialization.py --- # 浏览器初始化相关功能模块 import asyncio import os import time import json import logging from typing import Optional, Any, Dict, Tuple from playwright.async_api import Page as AsyncPage, Browser as AsyncBrowser, BrowserContext as AsyncBrowserContext, Error as PlaywrightAsyncError, expect as expect_async # 导入配置和模型 from config import * from models import ClientDisconnectedError logger = logging.getLogger("AIStudioProxyServer") async def _setup_network_interception_and_scripts(context: AsyncBrowserContext): """设置网络拦截和脚本注入""" try: from config.settings import ENABLE_SCRIPT_INJECTION if not ENABLE_SCRIPT_INJECTION: logger.info("脚本注入功能已禁用") return # 设置网络拦截 await _setup_model_list_interception(context) # 可选：仍然注入脚本作为备用方案 await _add_init_scripts_to_context(context) except Exception as e: logger.error(f"设置网络拦截和脚本注入时发生错误: {e}") async def _setup_model_list_interception(context: AsyncBrowserContext): """设置模型列表网络拦截""" try: async def handle_model_list_route(route): """处理模型列表请求的路由""" request = route.request # 检查是否是模型列表请求 if 'alkalimakersuite' in request.url and 'ListModels' in request.url: logger.info(f"🔍 拦截到模型列表请求: {request.url}") # 继续原始请求 response = await route.fetch() # 获取原始响应 original_body = await response.body() # 修改响应 modified_body = await _modify_model_list_response(original_body, request.url) # 返回修改后的响应 await route.fulfill( response=response, body=modified_body ) else: # 对于其他请求，直接继续 await route.continue_() # 注册路由拦截器 await context.route("**/*", handle_model_list_route) logger.info("✅ 已设置模型列表网络拦截") except Exception as e: logger.error(f"设置模型列表网络拦截时发生错误: {e}") async def _modify_model_list_response(original_body: bytes, url: str) -> bytes: """修改模型列表响应""" try: # 解码响应体 original_text = original_body.decode('utf-8') # 处理反劫持前缀 ANTI_HIJACK_PREFIX = ")]}'\n" has_prefix = False if original_text.startswith(ANTI_HIJACK_PREFIX): original_text = original_text[len(ANTI_HIJACK_PREFIX):] has_prefix = True # 解析JSON import json json_data = json.loads(original_text) # 注入模型 modified_data = await _inject_models_to_response(json_data, url) # 序列化回JSON modified_text = json.dumps(modified_data, separators=(',', ':')) # 重新添加前缀 if has_prefix: modified_text = ANTI_HIJACK_PREFIX + modified_text logger.info("✅ 成功修改模型列表响应") return modified_text.encode('utf-8') except Exception as e: logger.error(f"修改模型列表响应时发生错误: {e}") return original_body async def _inject_models_to_response(json_data: dict, url: str) -> dict: """向响应中注入模型""" try: from .operations import _get_injected_models # 获取要注入的模型 injected_models = _get_injected_models() if not injected_models: logger.info("没有要注入的模型") return json_data # 查找模型数组 models_array = _find_model_list_array(json_data) if not models_array: logger.warning("未找到模型数组结构") return json_data # 找到模板模型 template_model = _find_template_model(models_array) if not template_model: logger.warning("未找到模板模型") return json_data # 注入模型 for model in reversed(injected_models): # 反向以保持顺序 model_name = model['raw_model_path'] # 检查模型是否已存在 if not any(m[0] == model_name for m in models_array if isinstance(m, list) and len(m) > 0): # 创建新模型条目 new_model = json.loads(json.dumps(template_model)) # 深拷贝 new_model[0] = model_name # name new_model[3] = model['display_name'] # display name new_model[4] = model['description'] # description # 添加到开头 models_array.insert(0, new_model) logger.info(f"✅ 注入模型: {model['display_name']}") return json_data except Exception as e: logger.error(f"注入模型到响应时发生错误: {e}") return json_data def _find_model_list_array(obj): """递归查找模型列表数组""" if not obj: return None # 检查是否是模型数组 if isinstance(obj, list) and len(obj) > 0: if all(isinstance(item, list) and len(item) > 0 and isinstance(item[0], str) and item[0].startswith('models/') for item in obj): return obj # 递归搜索 if isinstance(obj, dict): for value in obj.values(): result = _find_model_list_array(value) if result: return result elif isinstance(obj, list): for item in obj: result = _find_model_list_array(item) if result: return result return None def _find_template_model(models_array): """查找模板模型""" if not models_array: return None # 寻找包含 'flash' 或 'pro' 的模型作为模板 for model in models_array: if isinstance(model, list) and len(model) > 7: model_name = model[0] if len(model) > 0 else "" if 'flash' in model_name.lower() or 'pro' in model_name.lower(): return model # 如果没找到，返回第一个有效模型 for model in models_array: if isinstance(model, list) and len(model) > 7: return model return None async def _add_init_scripts_to_context(context: AsyncBrowserContext): """在浏览器上下文中添加初始化脚本（备用方案）""" try: from config.settings import USERSCRIPT_PATH # 检查脚本文件是否存在 if not os.path.exists(USERSCRIPT_PATH): logger.info(f"脚本文件不存在，跳过脚本注入: {USERSCRIPT_PATH}") return # 读取脚本内容 with open(USERSCRIPT_PATH, 'r', encoding='utf-8') as f: script_content = f.read() # 清理UserScript头部 cleaned_script = _clean_userscript_headers(script_content) # 添加到上下文的初始化脚本 await context.add_init_script(cleaned_script) logger.info(f"✅ 已将脚本添加到浏览器上下文初始化脚本: {os.path.basename(USERSCRIPT_PATH)}") except Exception as e: logger.error(f"添加初始化脚本到上下文时发生错误: {e}") def _clean_userscript_headers(script_content: str) -> str: """清理UserScript头部信息""" lines = script_content.split('\n') cleaned_lines = [] in_userscript_block = False for line in lines: if line.strip().startswith('// ==UserScript=='): in_userscript_block = True continue elif line.strip().startswith('// ==/UserScript=='): in_userscript_block = False continue elif in_userscript_block: continue else: cleaned_lines.append(line) return '\n'.join(cleaned_lines) async def _initialize_page_logic(browser: AsyncBrowser): """初始化页面逻辑，连接到现有浏览器""" logger.info("--- 初始化页面逻辑 (连接到现有浏览器) ---") temp_context: Optional[AsyncBrowserContext] = None storage_state_path_to_use: Optional[str] = None launch_mode = os.environ.get('LAUNCH_MODE', 'debug') logger.info(f" 检测到启动模式: {launch_mode}") loop = asyncio.get_running_loop() if launch_mode == 'headless' or launch_mode == 'virtual_headless': auth_filename = os.environ.get('ACTIVE_AUTH_JSON_PATH') if auth_filename: constructed_path = auth_filename if os.path.exists(constructed_path): storage_state_path_to_use = constructed_path logger.info(f" 无头模式将使用的认证文件: {constructed_path}") else: logger.error(f"{launch_mode} 模式认证文件无效或不存在: '{constructed_path}'") raise RuntimeError(f"{launch_mode} 模式认证文件无效: '{constructed_path}'") else: logger.error(f"{launch_mode} 模式需要 ACTIVE_AUTH_JSON_PATH 环境变量，但未设置或为空。") raise RuntimeError(f"{launch_mode} 模式需要 ACTIVE_AUTH_JSON_PATH。") elif launch_mode == 'debug': logger.info(f" 调试模式: 尝试从环境变量 ACTIVE_AUTH_JSON_PATH 加载认证文件...") auth_filepath_from_env = os.environ.get('ACTIVE_AUTH_JSON_PATH') if auth_filepath_from_env and os.path.exists(auth_filepath_from_env): storage_state_path_to_use = auth_filepath_from_env logger.info(f" 调试模式将使用的认证文件 (来自环境变量): {storage_state_path_to_use}") elif auth_filepath_from_env: logger.warning(f" 调试模式下环境变量 ACTIVE_AUTH_JSON_PATH 指向的文件不存在: '{auth_filepath_from_env}'。不加载认证文件。") else: logger.info(" 调试模式下未通过环境变量提供认证文件。将使用浏览器当前状态。") elif launch_mode == "direct_debug_no_browser": logger.info(" direct_debug_no_browser 模式：不加载 storage_state，不进行浏览器操作。") else: logger.warning(f" ⚠️ 警告: 未知的启动模式 '{launch_mode}'。不加载 storage_state。") try: logger.info("创建新的浏览器上下文...") context_options: Dict[str, Any] = {'viewport': {'width': 460, 'height': 800}} if storage_state_path_to_use: context_options['storage_state'] = storage_state_path_to_use logger.info(f" (使用 storage_state='{os.path.basename(storage_state_path_to_use)}')") else: logger.info(" (不使用 storage_state)") # 代理设置需要从server模块中获取 import server if server.PLAYWRIGHT_PROXY_SETTINGS: context_options['proxy'] = server.PLAYWRIGHT_PROXY_SETTINGS logger.info(f" (浏览器上下文将使用代理: {server.PLAYWRIGHT_PROXY_SETTINGS['server']})") else: logger.info(" (浏览器上下文不使用显式代理配置)") context_options['ignore_https_errors'] = True logger.info(" (浏览器上下文将忽略 HTTPS 错误)") temp_context = await browser.new_context(**context_options) # 设置网络拦截和脚本注入 await _setup_network_interception_and_scripts(temp_context) found_page: Optional[AsyncPage] = None pages = temp_context.pages target_url_base = f"https://{AI_STUDIO_URL_PATTERN}" target_full_url = f"{target_url_base}prompts/new_chat" login_url_pattern = 'accounts.google.com' current_url = "" # 导入_handle_model_list_response - 需要延迟导入避免循环引用 from .operations import _handle_model_list_response for p_iter in pages: try: page_url_to_check = p_iter.url if not p_iter.is_closed() and target_url_base in page_url_to_check and "/prompts/" in page_url_to_check: found_page = p_iter current_url = page_url_to_check logger.info(f" 找到已打开的 AI Studio 页面: {current_url}") if found_page: logger.info(f" 为已存在的页面 {found_page.url} 添加模型列表响应监听器。") found_page.on("response", _handle_model_list_response) break except PlaywrightAsyncError as pw_err_url: logger.warning(f" 检查页面 URL 时出现 Playwright 错误: {pw_err_url}") except AttributeError as attr_err_url: logger.warning(f" 检查页面 URL 时出现属性错误: {attr_err_url}") except Exception as e_url_check: logger.warning(f" 检查页面 URL 时出现其他未预期错误: {e_url_check} (类型: {type(e_url_check).__name__})") if not found_page: logger.info(f"-> 未找到合适的现有页面，正在打开新页面并导航到 {target_full_url}...") found_page = await temp_context.new_page() if found_page: logger.info(f" 为新创建的页面添加模型列表响应监听器 (导航前)。") found_page.on("response", _handle_model_list_response) try: await found_page.goto(target_full_url, wait_until="domcontentloaded", timeout=90000) current_url = found_page.url logger.info(f"-> 新页面导航尝试完成。当前 URL: {current_url}") except Exception as new_page_nav_err: # 导入save_error_snapshot函数 from .operations import save_error_snapshot await save_error_snapshot("init_new_page_nav_fail") error_str = str(new_page_nav_err) if "NS_ERROR_NET_INTERRUPT" in error_str: logger.error("\n" + "="*30 + " 网络导航错误提示 " + "="*30) logger.error(f"❌ 导航到 '{target_full_url}' 失败，出现网络中断错误 (NS_ERROR_NET_INTERRUPT)。") logger.error(" 这通常表示浏览器在尝试加载页面时连接被意外断开。") logger.error(" 可能的原因及排查建议:") logger.error(" 1. 网络连接: 请检查你的本地网络连接是否稳定，并尝试在普通浏览器中访问目标网址。") logger.error(" 2. AI Studio 服务: 确认 aistudio.google.com 服务本身是否可用。") logger.error(" 3. 防火墙/代理/VPN: 检查本地防火墙、杀毒软件、代理或 VPN 设置。") logger.error(" 4. Camoufox 服务: 确认 launch_camoufox.py 脚本是否正常运行。") logger.error(" 5. 系统资源问题: 确保系统有足够的内存和 CPU 资源。") logger.error("="*74 + "\n") raise RuntimeError(f"导航新页面失败: {new_page_nav_err}") from new_page_nav_err if login_url_pattern in current_url: if launch_mode == 'headless': logger.error("无头模式下检测到重定向至登录页面，认证可能已失效。请更新认证文件。") raise RuntimeError("无头模式认证失败，需要更新认证文件。") else: print(f"\n{'='*20} 需要操作 {'='*20}", flush=True) login_prompt = " 检测到可能需要登录。如果浏览器显示登录页面，请在浏览器窗口中完成 Google 登录，然后在此处按 Enter 键继续..." print(USER_INPUT_START_MARKER_SERVER, flush=True) await loop.run_in_executor(None, input, login_prompt) print(USER_INPUT_END_MARKER_SERVER, flush=True) logger.info(" 用户已操作，正在检查登录状态...") try: await found_page.wait_for_url(f"**/{AI_STUDIO_URL_PATTERN}**", timeout=180000) current_url = found_page.url if login_url_pattern in current_url: logger.error("手动登录尝试后，页面似乎仍停留在登录页面。") raise RuntimeError("手动登录尝试后仍在登录页面。") logger.info(" ✅ 登录成功！请不要操作浏览器窗口，等待后续提示。") # 等待模型列表响应，确认登录成功 await _wait_for_model_list_and_handle_auth_save(temp_context, launch_mode, loop) except Exception as wait_login_err: from .operations import save_error_snapshot await save_error_snapshot("init_login_wait_fail") logger.error(f"登录提示后未能检测到 AI Studio URL 或保存状态时出错: {wait_login_err}", exc_info=True) raise RuntimeError(f"登录提示后未能检测到 AI Studio URL: {wait_login_err}") from wait_login_err elif target_url_base not in current_url or "/prompts/" not in current_url: from .operations import save_error_snapshot await save_error_snapshot("init_unexpected_page") logger.error(f"初始导航后页面 URL 意外: {current_url}。期望包含 '{target_url_base}' 和 '/prompts/'。") raise RuntimeError(f"初始导航后出现意外页面: {current_url}。") logger.info(f"-> 确认当前位于 AI Studio 对话页面: {current_url}") await found_page.bring_to_front() try: input_wrapper_locator = found_page.locator('ms-prompt-input-wrapper') await expect_async(input_wrapper_locator).to_be_visible(timeout=35000) await expect_async(found_page.locator(INPUT_SELECTOR)).to_be_visible(timeout=10000) logger.info("-> ✅ 核心输入区域可见。") model_name_locator = found_page.locator('mat-select[data-test-ms-model-selector] div.model-option-content span.gmat-body-medium') try: model_name_on_page = await model_name_locator.first.inner_text(timeout=5000) logger.info(f"-> 🤖 页面检测到的当前模型: {model_name_on_page}") except PlaywrightAsyncError as e: logger.error(f"获取模型名称时出错 (model_name_locator): {e}") raise result_page_instance = found_page result_page_ready = True # 脚本注入已在上下文创建时完成，无需在此处重复注入 logger.info(f"✅ 页面逻辑初始化成功。") return result_page_instance, result_page_ready except Exception as input_visible_err: from .operations import save_error_snapshot await save_error_snapshot("init_fail_input_timeout") logger.error(f"页面初始化失败：核心输入区域未在预期时间内变为可见。最后的 URL 是 {found_page.url}", exc_info=True) raise RuntimeError(f"页面初始化失败：核心输入区域未在预期时间内变为可见。最后的 URL 是 {found_page.url}") from input_visible_err except Exception as e_init_page: logger.critical(f"❌ 页面逻辑初始化期间发生严重意外错误: {e_init_page}", exc_info=True) if temp_context: try: logger.info(f" 尝试关闭临时的浏览器上下文 due to initialization error.") await temp_context.close() logger.info(" ✅ 临时浏览器上下文已关闭。") except Exception as close_err: logger.warning(f" ⚠️ 关闭临时浏览器上下文时出错: {close_err}") from .operations import save_error_snapshot await save_error_snapshot("init_unexpected_error") raise RuntimeError(f"页面初始化意外错误: {e_init_page}") from e_init_page async def _close_page_logic(): """关闭页面逻辑""" # 需要访问全局变量 import server logger.info("--- 运行页面逻辑关闭 --- ") if server.page_instance and not server.page_instance.is_closed(): try: await server.page_instance.close() logger.info(" ✅ 页面已关闭") except PlaywrightAsyncError as pw_err: logger.warning(f" ⚠️ 关闭页面时出现Playwright错误: {pw_err}") except asyncio.TimeoutError as timeout_err: logger.warning(f" ⚠️ 关闭页面时超时: {timeout_err}") except Exception as other_err: logger.error(f" ⚠️ 关闭页面时出现意外错误: {other_err} (类型: {type(other_err).__name__})", exc_info=True) server.page_instance = None server.is_page_ready = False logger.info("页面逻辑状态已重置。") return None, False async def signal_camoufox_shutdown(): """发送关闭信号到Camoufox服务器""" logger.info(" 尝试发送关闭信号到 Camoufox 服务器 (此功能可能已由父进程处理)...") ws_endpoint = os.environ.get('CAMOUFOX_WS_ENDPOINT') if not ws_endpoint: logger.warning(" ⚠️ 无法发送关闭信号：未找到 CAMOUFOX_WS_ENDPOINT 环境变量。") return # 需要访问全局浏览器实例 import server if not server.browser_instance or not server.browser_instance.is_connected(): logger.warning(" ⚠️ 浏览器实例已断开或未初始化，跳过关闭信号发送。") return try: await asyncio.sleep(0.2) logger.info(" ✅ (模拟) 关闭信号已处理。") except Exception as e: logger.error(f" ⚠️ 发送关闭信号过程中捕获异常: {e}", exc_info=True) async def _wait_for_model_list_and_handle_auth_save(temp_context, launch_mode, loop): """等待模型列表响应并处理认证保存""" import server # 等待模型列表响应，确认登录成功 logger.info(" 等待模型列表响应以确认登录成功...") try: # 等待模型列表事件，最多等待30秒 await asyncio.wait_for(server.model_list_fetch_event.wait(), timeout=30.0) logger.info(" ✅ 检测到模型列表响应，登录确认成功！") except asyncio.TimeoutError: logger.warning(" ⚠️ 等待模型列表响应超时，但继续处理认证保存...") # 检查是否启用自动确认 if AUTO_CONFIRM_LOGIN: print("\n" + "="*50, flush=True) print(" ✅ 登录成功！检测到模型列表响应。", flush=True) print(" 🤖 自动确认模式已启用，将自动保存认证状态...", flush=True) # 自动保存认证状态 await _handle_auth_file_save_auto(temp_context) print("="*50 + "\n", flush=True) return # 手动确认模式 print("\n" + "="*50, flush=True) print(" 【用户交互】需要您的输入!", flush=True) print(" ✅ 登录成功！检测到模型列表响应。", flush=True) should_save_auth_choice = '' if AUTO_SAVE_AUTH and launch_mode == 'debug': logger.info(" 自动保存认证模式已启用，将自动保存认证状态...") should_save_auth_choice = 'y' else: save_auth_prompt = " 是否要将当前的浏览器认证状态保存到文件？ (y/N): " print(USER_INPUT_START_MARKER_SERVER, flush=True) try: auth_save_input_future = loop.run_in_executor(None, input, save_auth_prompt) should_save_auth_choice = await asyncio.wait_for(auth_save_input_future, timeout=AUTH_SAVE_TIMEOUT) except asyncio.TimeoutError: print(f" 输入等待超时({AUTH_SAVE_TIMEOUT}秒)。默认不保存认证状态。", flush=True) should_save_auth_choice = 'n' finally: print(USER_INPUT_END_MARKER_SERVER, flush=True) if should_save_auth_choice.strip().lower() == 'y': await _handle_auth_file_save(temp_context, loop) else: print(" 好的，不保存认证状态。", flush=True) print("="*50 + "\n", flush=True) async def _handle_auth_file_save(temp_context, loop): """处理认证文件保存（手动模式）""" os.makedirs(SAVED_AUTH_DIR, exist_ok=True) default_auth_filename = f"auth_state_{int(time.time())}.json" print(USER_INPUT_START_MARKER_SERVER, flush=True) filename_prompt_str = f" 请输入保存的文件名 (默认为: {default_auth_filename}，输入 'cancel' 取消保存): " chosen_auth_filename = '' try: filename_input_future = loop.run_in_executor(None, input, filename_prompt_str) chosen_auth_filename = await asyncio.wait_for(filename_input_future, timeout=AUTH_SAVE_TIMEOUT) except asyncio.TimeoutError: print(f" 输入文件名等待超时({AUTH_SAVE_TIMEOUT}秒)。将使用默认文件名: {default_auth_filename}", flush=True) chosen_auth_filename = default_auth_filename finally: print(USER_INPUT_END_MARKER_SERVER, flush=True) # 检查用户是否选择取消 if chosen_auth_filename.strip().lower() == 'cancel': print(" 用户选择取消保存认证状态。", flush=True) return final_auth_filename = chosen_auth_filename.strip() or default_auth_filename if not final_auth_filename.endswith(".json"): final_auth_filename += ".json" auth_save_path = os.path.join(SAVED_AUTH_DIR, final_auth_filename) try: await temp_context.storage_state(path=auth_save_path) print(f" ✅ 认证状态已成功保存到: {auth_save_path}", flush=True) except Exception as save_state_err: logger.error(f" ❌ 保存认证状态失败: {save_state_err}", exc_info=True) print(f" ❌ 保存认证状态失败: {save_state_err}", flush=True) async def _handle_auth_file_save_auto(temp_context): """处理认证文件保存（自动模式）""" os.makedirs(SAVED_AUTH_DIR, exist_ok=True) # 生成基于时间戳的文件名 timestamp = int(time.time()) auto_auth_filename = f"auth_auto_{timestamp}.json" auth_save_path = os.path.join(SAVED_AUTH_DIR, auto_auth_filename) try: await temp_context.storage_state(path=auth_save_path) print(f" ✅ 认证状态已自动保存到: {auth_save_path}", flush=True) logger.info(f" 自动保存认证状态成功: {auth_save_path}") except Exception as save_state_err: logger.error(f" ❌ 自动保存认证状态失败: {save_state_err}", exc_info=True) print(f" ❌ 自动保存认证状态失败: {save_state_err}", flush=True)