#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 金山文档上传模块 - 精简版 使用Playwright自动化上传截图到金山文档表格 移除了队列、并发控制,改为单任务顺序执行 """ import base64 import os import re import time from io import BytesIO from typing import Any, Dict, Optional, Callable from urllib.parse import urlparse try: from playwright.sync_api import sync_playwright, TimeoutError as PlaywrightTimeoutError except ImportError: sync_playwright = None class PlaywrightTimeoutError(Exception): pass class KDocsUploader: """金山文档上传器""" def __init__(self, log_callback: Optional[Callable] = None): self._playwright = None self._browser = None self._context = None self._page = None self._doc_url: Optional[str] = None self._last_error: Optional[str] = None self._logged_in = False self._log_callback = log_callback def log(self, msg: str): """记录日志""" if self._log_callback: self._log_callback(msg) def _ensure_playwright(self, use_storage_state: bool = True) -> bool: """确保Playwright已启动""" if sync_playwright is None: self._last_error = "playwright 未安装" return False try: from config import KDOCS_LOGIN_STATE_FILE if self._playwright is None: self._playwright = sync_playwright().start() if self._browser is None: # 调试模式:有头模式,方便查看浏览器行为 # 生产环境改回 "true" headless = os.environ.get("KDOCS_HEADLESS", "false").lower() != "false" # 使用系统安装的Chrome浏览器(支持微信快捷登录) # channel='chrome' 会使用系统Chrome,而不是Playwright自带的Chromium chrome_args = [ "--disable-blink-features=AutomationControlled", # 隐藏自动化特征 "--disable-features=DialMediaRouteProvider", # 禁用本地网络发现提示 "--allow-running-insecure-content", ] try: self._browser = self._playwright.chromium.launch( headless=headless, channel='chrome', # 使用系统Chrome args=chrome_args ) self.log("[KDocs] 使用系统Chrome浏览器") except Exception as e: # 如果系统没有Chrome,回退到Chromium self.log(f"[KDocs] 系统Chrome不可用({e}),使用Chromium") self._browser = self._playwright.chromium.launch(headless=headless, args=chrome_args) if self._context is None: storage_state = str(KDOCS_LOGIN_STATE_FILE) # 创建context时的通用配置 context_options = { "permissions": ["clipboard-read", "clipboard-write"], # 剪贴板权限 "ignore_https_errors": True, } if use_storage_state and os.path.exists(storage_state): context_options["storage_state"] = storage_state self._context = self._browser.new_context(**context_options) # 授予本地网络访问权限(用于微信快捷登录检测) try: self._context.grant_permissions( ["clipboard-read", "clipboard-write"], origin="https://account.wps.cn" ) except Exception: pass if self._page is None or self._page.is_closed(): self._page = self._context.new_page() self._page.set_default_timeout(60000) return True except Exception as e: self._last_error = f"浏览器启动失败: {e}" self._cleanup_browser() return False def _cleanup_browser(self): """清理浏览器资源""" for attr in ['_page', '_context', '_browser', '_playwright']: obj = getattr(self, attr, None) if obj: try: if hasattr(obj, 'close'): obj.close() elif hasattr(obj, 'stop'): obj.stop() except Exception: pass setattr(self, attr, None) def _open_document(self, doc_url: str) -> bool: """打开金山文档""" try: self._doc_url = doc_url self._ensure_clipboard_permissions(doc_url) self._page.goto(doc_url, wait_until="domcontentloaded", timeout=30000) time.sleep(3) # 等待页面完全加载,包括登录按钮 return True except Exception as e: self._last_error = f"打开文档失败: {e}" return False def _ensure_clipboard_permissions(self, doc_url: str): """授予剪贴板权限""" if not self._context or not doc_url: return try: parsed = urlparse(doc_url) if not parsed.scheme or not parsed.netloc: return origin = f"{parsed.scheme}://{parsed.netloc}" self._context.grant_permissions(["clipboard-read", "clipboard-write"], origin=origin) except Exception: pass def _is_login_url(self, url: str) -> bool: """检查是否是登录页面""" if not url: return False lower = url.lower() if "account.wps.cn" in lower or "passport" in lower: return True if "login" in lower and "kdocs.cn" not in lower: return True return False def _page_has_login_gate(self, page) -> bool: """检查页面是否需要登录""" url = getattr(page, "url", "") or "" # 如果URL已经是文档页面,说明已登录成功 if "kdocs.cn/l/" in url or "www.kdocs.cn/l/" in url: # 但可能有邀请对话框,先尝试点击关闭 try: join_btn = page.get_by_role("button", name="登录并加入编辑") if join_btn.count() > 0 and join_btn.first.is_visible(timeout=500): self.log("[KDocs] 点击加入编辑按钮") join_btn.first.click() time.sleep(1) except Exception: pass # 已经在文档页面,算作已登录 return False # 检查是否在登录页面 if self._is_login_url(url): self.log(f"[KDocs] 检测到登录页面URL: {url}") return True # 只检查登录页面上的登录按钮(排除文档页面的邀请对话框) login_buttons = ["立即登录", "去登录"] for text in login_buttons: try: btn = page.get_by_role("button", name=text) if btn.count() > 0 and btn.first.is_visible(timeout=500): self.log(f"[KDocs] 检测到登录按钮: {text}") return True except Exception: pass # 检查是否有二维码元素可见(说明还在等待扫码) try: qr_selectors = ["canvas", "img[class*='qr']", "div[class*='qrcode']"] for selector in qr_selectors: qr = page.locator(selector) if qr.count() > 0: for i in range(min(qr.count(), 3)): el = qr.nth(i) try: if el.is_visible(timeout=200): box = el.bounding_box() if box and 80 <= box.get("width", 0) <= 400: self.log(f"[KDocs] 检测到二维码元素: {selector}") return True except Exception: pass except Exception: pass return False def _is_logged_in(self) -> bool: """检查是否已登录""" if not self._page or self._page.is_closed(): return False return not self._page_has_login_gate(self._page) def _save_login_state(self): """保存登录状态""" try: from config import KDOCS_LOGIN_STATE_FILE storage_state = str(KDOCS_LOGIN_STATE_FILE) KDOCS_LOGIN_STATE_FILE.parent.mkdir(parents=True, exist_ok=True) self._context.storage_state(path=storage_state) self.log("[KDocs] 登录状态已保存") except Exception as e: self.log(f"[KDocs] 保存登录状态失败: {e}") def _ensure_login_dialog(self, use_quick_login: bool = False): """确保打开登录对话框 Args: use_quick_login: 是否尝试使用微信快捷登录 """ agree_names = ["同意", "同意并继续", "我同意", "确定", "确认"] # 循环处理登录流程 max_clicks = 8 for round_num in range(max_clicks): clicked = False current_url = self._page.url # 检查是否已经到达文档页面(登录成功) # 需要确保不是临时跳转,等待页面稳定 if "kdocs.cn/l/" in current_url or "www.kdocs.cn/l/" in current_url: time.sleep(1) # 等待页面稳定 stable_url = self._page.url if "kdocs.cn/l/" in stable_url and "account.wps.cn" not in stable_url: self.log("[KDocs] 已到达文档页面,登录成功") return # 1. 先检查是否有隐私协议同意按钮 for name in agree_names: try: btn = self._page.get_by_role("button", name=name) if btn.count() > 0 and btn.first.is_visible(timeout=300): self.log(f"[KDocs] 点击同意按钮: {name}") btn.first.click() time.sleep(1) clicked = True break except Exception: pass if clicked: continue # 2. 如果启用快捷登录且在登录页面(account.wps.cn),尝试点击"微信快捷登录" if use_quick_login and "account.wps.cn" in current_url: try: quick_login = self._page.get_by_text("微信快捷登录", exact=False) if quick_login.count() > 0 and quick_login.first.is_visible(timeout=500): self.log("[KDocs] 点击微信快捷登录") quick_login.first.click() time.sleep(3) # 等待快捷登录处理 # 检查是否登录成功 if "kdocs.cn/l/" in self._page.url: self.log("[KDocs] 微信快捷登录成功") return clicked = True continue except Exception: pass # 3. 点击"立即登录"进入登录页面 try: btn = self._page.get_by_role("button", name="立即登录") if btn.count() > 0 and btn.first.is_visible(timeout=500): self.log("[KDocs] 点击立即登录") btn.first.click() time.sleep(2) clicked = True continue except Exception: pass # 4. 点击"登录并加入编辑"(文档页面的邀请对话框) try: btn = self._page.get_by_role("button", name="登录并加入编辑") if btn.count() > 0 and btn.first.is_visible(timeout=500): self.log("[KDocs] 点击登录并加入编辑") btn.first.click() time.sleep(1.5) clicked = True continue except Exception: pass # 如果没有点击到任何按钮,退出循环 if not clicked: self.log("[KDocs] 未找到更多可点击的按钮") break # 最后确保点击微信扫码登录(切换到扫码模式) wechat_names = ["微信登录", "微信扫码登录", "扫码登录", "微信扫码"] for name in wechat_names: try: btn = self._page.get_by_role("button", name=name) if btn.is_visible(timeout=1000): self.log(f"[KDocs] 点击微信登录: {name}") btn.click() time.sleep(1) return except Exception: pass # 尝试用文本查找微信登录 for name in wechat_names: try: el = self._page.get_by_text(name, exact=False).first if el.is_visible(timeout=500): self.log(f"[KDocs] 点击微信登录文本: {name}") el.click() time.sleep(1) return except Exception: pass self.log("[KDocs] 未找到登录按钮,可能页面已在登录状态或需要手动操作") def _capture_qr_image(self) -> Optional[bytes]: """捕获登录二维码图片""" # 查找二维码元素的选择器 selectors = [ "canvas", "img[src*='qr']", "img[class*='qr']", "img[class*='code']", "div[class*='qr'] img", "div[class*='qrcode'] img", "div[class*='scan'] img", ".qrcode img", ".qr-code img", "img", # 最后尝试所有图片 ] # 先在主页面查找 for selector in selectors: result = self._try_capture_qr_with_selector(self._page, selector) if result: return result # 尝试在iframe中查找 try: frames = self._page.frames for frame in frames: if frame == self._page.main_frame: continue for selector in selectors[:5]: # 只用前几个选择器 result = self._try_capture_qr_with_selector(frame, selector) if result: return result except Exception: pass return None def _try_capture_qr_with_selector(self, page_or_frame, selector: str) -> Optional[bytes]: """尝试用指定选择器捕获二维码""" try: locator = page_or_frame.locator(selector) count = locator.count() for i in range(min(count, 10)): el = locator.nth(i) try: if not el.is_visible(timeout=300): continue box = el.bounding_box() if not box: continue w, h = box.get("width", 0), box.get("height", 0) # 二维码通常是正方形,大小在100-400之间 if 80 <= w <= 400 and 80 <= h <= 400 and abs(w - h) < 60: screenshot = el.screenshot() if screenshot and len(screenshot) > 500: return screenshot except Exception: continue except Exception: pass return None def request_qr(self, force: bool = False) -> Dict[str, Any]: """ 请求登录二维码 Args: force: 是否强制重新登录 Returns: { "success": bool, "logged_in": bool, # 是否已登录 "qr_image": str, # base64编码的二维码图片 "error": str # 错误信息 } """ from config import get_config, KDOCS_LOGIN_STATE_FILE config = get_config() doc_url = config.kdocs.doc_url.strip() if not doc_url: return {"success": False, "error": "未配置金山文档链接"} if force: # 清除登录状态 try: if KDOCS_LOGIN_STATE_FILE.exists(): KDOCS_LOGIN_STATE_FILE.unlink() except Exception: pass self._cleanup_browser() if not self._ensure_playwright(use_storage_state=not force): return {"success": False, "error": self._last_error or "浏览器不可用"} if not self._open_document(doc_url): return {"success": False, "error": self._last_error or "打开文档失败"} # 检查是否已登录 self.log(f"[KDocs] 当前页面URL: {self._page.url}") if not force and self._is_logged_in(): self._logged_in = True self._save_login_state() return {"success": True, "logged_in": True, "qr_image": ""} # 需要登录,获取二维码 self.log("[KDocs] 需要登录,尝试打开登录对话框...") self._ensure_login_dialog() time.sleep(2) # 等待登录对话框加载 self.log("[KDocs] 尝试捕获二维码...") qr_image = None for i in range(15): # 增加尝试次数 qr_image = self._capture_qr_image() if qr_image and len(qr_image) > 1024: self.log(f"[KDocs] 二维码捕获成功,大小: {len(qr_image)} bytes") break self.log(f"[KDocs] 第{i+1}次尝试捕获二维码...") time.sleep(1) if not qr_image: # 尝试截取整个页面帮助调试 self.log("[KDocs] 二维码捕获失败,当前页面可能没有显示二维码") return {"success": False, "error": "二维码获取失败,请检查网络或手动打开金山文档链接确认"} return { "success": True, "logged_in": False, "qr_image": base64.b64encode(qr_image).decode("ascii"), } def check_login_status(self) -> Dict[str, Any]: """检查登录状态(不重新打开页面,只检查当前状态)""" # 如果页面不存在或已关闭,说明还没开始登录流程 if not self._page or self._page.is_closed(): return {"success": False, "logged_in": False, "error": "页面未打开"} try: clicked_confirm = False # 在主页面和所有iframe中查找确认按钮 frames_to_check = [self._page] + list(self._page.frames) for frame in frames_to_check: if clicked_confirm: break # 尝试点击确认登录按钮(微信扫码后PC端需要再点一下确认) confirm_names = ["确认登录", "确定登录", "登录", "确定", "确认", "同意并登录"] for name in confirm_names: try: confirm_btn = frame.get_by_role("button", name=name) if confirm_btn.count() > 0 and confirm_btn.first.is_visible(timeout=200): self.log(f"[KDocs] 找到确认按钮: {name}") confirm_btn.first.click() clicked_confirm = True time.sleep(3) break except Exception: pass # 如果按钮角色没找到,尝试用文本查找 if not clicked_confirm: for name in confirm_names: try: el = frame.get_by_text(name, exact=True) if el.count() > 0 and el.first.is_visible(timeout=200): self.log(f"[KDocs] 找到确认文本: {name}") el.first.click() clicked_confirm = True time.sleep(3) break except Exception: pass # 尝试用CSS选择器查找 if not clicked_confirm: try: # WPS登录页面的确认按钮可能的选择器 selectors = [ "button.ant-btn-primary", "button[type='primary']", ".confirm-btn", ".login-confirm", ".btn-primary", ".wps-btn-primary", "a.confirm", "div.confirm", "[class*='confirm']", "[class*='login-btn']" ] for selector in selectors: btns = frame.locator(selector) if btns.count() > 0: for i in range(min(btns.count(), 3)): btn = btns.nth(i) try: if btn.is_visible(timeout=100): btn_text = btn.inner_text() or "" if any(kw in btn_text for kw in ["确认", "登录", "确定"]): self.log(f"[KDocs] 找到按钮(CSS): {btn_text}") btn.click() clicked_confirm = True time.sleep(3) break except Exception: pass if clicked_confirm: break except Exception: pass # 如果点击了确认按钮,等待页面自动跳转(不要reload!) if clicked_confirm: self.log("[KDocs] 已点击确认,等待页面跳转...") time.sleep(3) # 等待页面自动跳转 # 检查当前URL是否已经到达文档页面 current_url = self._page.url self.log(f"[KDocs] 当前URL: {current_url}") # 直接检查URL判断是否已登录 if "kdocs.cn/l/" in current_url and "account.wps.cn" not in current_url: # 已到达文档页面,登录成功 logged_in = True # 尝试点击可能存在的"加入编辑"按钮 try: join_btn = self._page.get_by_role("button", name="登录并加入编辑") if join_btn.count() > 0 and join_btn.first.is_visible(timeout=500): self.log("[KDocs] 点击加入编辑") join_btn.first.click() time.sleep(1) except Exception: pass else: # 还在登录页面或其他页面 logged_in = self._is_logged_in() self._logged_in = logged_in if logged_in: self._save_login_state() self.log("[KDocs] 登录状态检测:已登录") return {"success": True, "logged_in": logged_in} except Exception as e: return {"success": False, "logged_in": False, "error": str(e)} def _navigate_to_cell(self, cell_address: str): """导航到指定单元格""" try: name_box = self._page.locator("input.edit-box").first name_box.click() name_box.fill(cell_address) name_box.press("Enter") except Exception: name_box = self._page.locator('#root input[type="text"]').first name_box.click() name_box.fill(cell_address) name_box.press("Enter") time.sleep(0.3) def _get_current_cell_address(self) -> str: """获取当前单元格地址""" try: name_box = self._page.locator("input.edit-box").first value = name_box.input_value() if value and re.match(r"^[A-Z]+\d+$", value.upper()): return value.upper() except Exception: pass return "" def _search_and_get_row(self, search_text: str, expected_col: str = None, row_start: int = 0, row_end: int = 0) -> int: """搜索并获取行号""" # 打开搜索 self._page.keyboard.press("Control+f") time.sleep(0.3) # 输入搜索内容 try: search_input = self._page.get_by_role("textbox").nth(3) if search_input.is_visible(timeout=500): search_input.fill(search_text) except Exception: pass time.sleep(0.2) # 点击查找 try: find_btn = self._page.get_by_role("button", name="查找").first find_btn.click() except Exception: self._page.keyboard.press("Enter") time.sleep(0.3) # 获取当前位置 self._page.keyboard.press("Escape") time.sleep(0.3) address = self._get_current_cell_address() if not address: return -1 # 提取行号 match = re.search(r"(\d+)$", address) if not match: return -1 row_num = int(match.group(1)) col_letter = "".join(c for c in address if c.isalpha()).upper() # 检查列 if expected_col and col_letter != expected_col.upper(): return -1 # 检查行范围 if row_start > 0 and row_num < row_start: return -1 if row_end > 0 and row_num > row_end: return -1 return row_num def _upload_image_to_cell(self, row_num: int, image_path: str, image_col: str) -> bool: """上传图片到单元格""" cell_address = f"{image_col}{row_num}" self._navigate_to_cell(cell_address) time.sleep(0.3) # 清除单元格内容 try: self._page.keyboard.press("Escape") time.sleep(0.2) self._page.keyboard.press("Delete") time.sleep(0.3) except Exception: pass # 插入 -> 图片 -> 单元格图片 try: insert_btn = self._page.get_by_role("button", name="插入") insert_btn.click() time.sleep(0.3) image_btn = self._page.get_by_role("button", name="图片") image_btn.click() time.sleep(0.3) cell_image_option = self._page.get_by_role("option", name="单元格图片") cell_image_option.click() time.sleep(0.2) local_option = self._page.get_by_role("option", name="本地") with self._page.expect_file_chooser() as fc_info: local_option.click() file_chooser = fc_info.value file_chooser.set_files(image_path) time.sleep(2) self.log(f"[KDocs] 图片已上传到 {cell_address}") return True except Exception as e: self._last_error = f"上传图片失败: {e}" return False def upload_image( self, image_path: str, unit: str, name: str, ) -> Dict[str, Any]: """ 上传截图到金山文档 Args: image_path: 图片路径 unit: 县区名(用于定位行) name: 姓名(用于定位行) Returns: {"success": bool, "error": str} """ from config import get_config config = get_config() kdocs_config = config.kdocs if not kdocs_config.enabled: return {"success": False, "error": "金山文档上传未启用"} doc_url = kdocs_config.doc_url.strip() if not doc_url: return {"success": False, "error": "未配置金山文档链接"} if not unit or not name: return {"success": False, "error": "缺少县区或姓名"} if not image_path or not os.path.exists(image_path): return {"success": False, "error": "图片文件不存在"} if not self._ensure_playwright(): return {"success": False, "error": self._last_error or "浏览器不可用"} if not self._open_document(doc_url): return {"success": False, "error": self._last_error or "打开文档失败"} if not self._is_logged_in(): return {"success": False, "error": "未登录,请先扫码登录"} try: # 选择工作表 if kdocs_config.sheet_name: try: tab = self._page.locator("[role='tab']").filter(has_text=kdocs_config.sheet_name) if tab.count() > 0: tab.first.click() time.sleep(0.5) except Exception: pass # 搜索姓名找到行 self.log(f"[KDocs] 搜索人员: {name}") row_num = self._search_and_get_row( name, expected_col=kdocs_config.name_column, row_start=kdocs_config.row_start, row_end=kdocs_config.row_end, ) if row_num < 0: return {"success": False, "error": f"未找到人员: {name}"} self.log(f"[KDocs] 找到人员在第 {row_num} 行") # 上传图片 if self._upload_image_to_cell(row_num, image_path, kdocs_config.image_column): return {"success": True} else: return {"success": False, "error": self._last_error or "上传失败"} except Exception as e: return {"success": False, "error": str(e)} def clear_login(self): """清除登录状态""" from config import KDOCS_LOGIN_STATE_FILE try: if KDOCS_LOGIN_STATE_FILE.exists(): KDOCS_LOGIN_STATE_FILE.unlink() except Exception: pass self._logged_in = False self._cleanup_browser() def close(self): """关闭上传器""" self._cleanup_browser() def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): self.close() return False # 全局实例 _uploader: Optional[KDocsUploader] = None def get_kdocs_uploader() -> KDocsUploader: """获取金山文档上传器实例""" global _uploader if _uploader is None: _uploader = KDocsUploader() return _uploader