Files
zsglpt-pc/core/kdocs_uploader.py
237899745 9743186a9e feat: 添加依赖自动检测与安装、选项记忆、KDocs登录优化
- 新增依赖检测模块:启动时自动检测wkhtmltoimage和Playwright Chromium
- 新增依赖安装对话框:缺失时提示用户一键下载安装
- 修复选项记忆功能:浏览类型、自动截图、自动上传选项现在会保存
- 优化KDocs登录检测:未登录时自动切换到金山文档页面并显示二维码
- 简化日志输出:移除debug信息,保留用户友好的状态提示
- 新增账号变化信号:账号管理页面的修改会自动同步到浏览任务页面

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
2026-01-19 01:28:06 +08:00

767 lines
27 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/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 _find_visible_element(self, text: str, use_role: bool = False, role: str = "button"):
"""找到包含指定文本的可见元素
Args:
text: 要查找的文本
use_role: 是否使用role查找默认False使用文本查找
role: 使用role查找时的角色类型
Returns:
可见的元素或None
"""
if not self._page:
return None
try:
if use_role:
els = self._page.get_by_role(role, name=text)
else:
els = self._page.locator(f"text={text}")
count = els.count()
for i in range(count):
el = els.nth(i)
try:
if el.is_visible(timeout=500):
box = el.bounding_box()
if box and box.get('width', 0) > 0 and box.get('height', 0) > 0:
return el
except Exception:
pass
except Exception:
pass
return None
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:
# 默认无头模式,设置环境变量 KDOCS_HEADLESS=false 可切换为有头模式调试
headless = os.environ.get("KDOCS_HEADLESS", "true").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
)
except Exception as e:
# 如果系统没有Chrome回退到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):
join_btn.first.click()
time.sleep(1)
except Exception:
pass
# 已经在文档页面,算作已登录
return False
# 检查是否在登录页面
if self._is_login_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):
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:
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)
except Exception:
pass
def _ensure_login_dialog(self, use_quick_login: bool = False):
"""确保打开登录对话框并进入扫码页面"""
buttons_priority = [
"登录并加入编辑",
"立即登录",
"去登录",
]
max_clicks = 12
for _ in range(max_clicks):
clicked = False
current_url = self._page.url
# 检查是否已经到达文档页面(登录成功)
if "kdocs.cn/l/" in current_url and "account.wps.cn" not in current_url:
time.sleep(1)
has_login_elements = (
self._find_visible_element("立即登录", use_role=True) or
self._find_visible_element("登录并加入编辑", use_role=True) or
self._find_visible_element("微信扫码登录") or
self._find_visible_element("微信快捷登录")
)
if not has_login_elements:
return
# 检查是否已经到达登录二维码页面
qr_page_indicators = ["微信扫码登录", "微信快捷登录"]
for indicator in qr_page_indicators:
if self._find_visible_element(indicator):
return
# 按优先级点击登录按钮
for btn_name in buttons_priority:
el = self._find_visible_element(btn_name, use_role=True)
if el:
el.click(force=True)
time.sleep(2)
clicked = True
break
if not clicked:
for btn_name in buttons_priority:
el = self._find_visible_element(btn_name)
if el:
el.click(force=True)
time.sleep(2)
clicked = True
break
if not clicked:
time.sleep(1)
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 "打开文档失败"}
# 检查是否已登录
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._ensure_login_dialog()
time.sleep(2)
qr_image = None
for _ in range(15):
qr_image = self._capture_qr_image()
if qr_image and len(qr_image) > 1024:
break
time.sleep(1)
if not qr_image:
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):
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):
el.first.click()
clicked_confirm = True
time.sleep(3)
break
except Exception:
pass
# 尝试用CSS选择器查找
if not clicked_confirm:
try:
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 ["确认", "登录", "确定"]):
btn.click()
clicked_confirm = True
time.sleep(3)
break
except Exception:
pass
if clicked_confirm:
break
except Exception:
pass
if clicked_confirm:
time.sleep(3)
current_url = self._page.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):
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()
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)
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
# 搜索姓名找到行
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}"}
# 上传图片
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