234 lines
10 KiB
Python
234 lines
10 KiB
Python
"""
|
||
ERP 登录模块 - DrissionPage
|
||
"""
|
||
|
||
import os
|
||
import time
|
||
import datetime
|
||
import tempfile
|
||
from pathlib import Path
|
||
from dotenv import load_dotenv
|
||
from DrissionPage import ChromiumPage, ChromiumOptions
|
||
|
||
# ── 加载 .env ─────────────────────────────────────────────────────────────────
|
||
load_dotenv(Path(__file__).parent / ".env")
|
||
|
||
# 强制 Python 的 websocket 客户端忽略本地代理，防止出现 Handshake status 404 Not Found
|
||
os.environ["NO_PROXY"] = "localhost,127.0.0.1,::1"
|
||
os.environ["no_proxy"] = "localhost,127.0.0.1,::1"
|
||
|
||
ERP_URL = os.getenv("ERP_URL", "https://yunmes.tftykj.cn/#")
|
||
ERP_TENANT = os.getenv("ERP_TENANT", "")
|
||
ERP_USERNAME = os.getenv("ERP_USERNAME", "")
|
||
ERP_PASSWORD = os.getenv("ERP_PASSWORD", "")
|
||
|
||
|
||
def is_docker_env() -> bool:
|
||
"""判断当前是否运行在 Docker 容器中。"""
|
||
return os.path.exists("/.dockerenv")
|
||
|
||
|
||
def get_docker_user_data_dir() -> Path:
|
||
"""
|
||
为每个进程分配独立的 Chromium 用户目录。
|
||
多 worker 或异常重启后复用同一 profile，容易导致 DevTools 会话错乱。
|
||
"""
|
||
base_dir = Path(os.getenv("DRISSION_USER_DATA_ROOT", tempfile.gettempdir())) / "drission_profiles"
|
||
user_data_dir = base_dir / f"worker_{os.getpid()}"
|
||
user_data_dir.mkdir(parents=True, exist_ok=True)
|
||
return user_data_dir
|
||
|
||
|
||
# ── 日志 ──────────────────────────────────────────────────────────────────────
|
||
def log(level: str, msg: str):
|
||
icons = {"INFO": "ℹ️ ", "OK": "✅", "WARN": "⚠️ ", "ERR": "❌"}
|
||
ts = datetime.datetime.now().strftime("%H:%M:%S.%f")[:-3]
|
||
print(f"[{ts}] {icons.get(level, ' ')} [{level}] {msg}", flush=True)
|
||
|
||
|
||
def dump_page_state(page: ChromiumPage, label: str = ""):
|
||
"""仅在出错时调用，打印页面快照用于排查。"""
|
||
tag = f" ({label})" if label else ""
|
||
log("WARN", f"─── 页面快照{tag} ───")
|
||
log("WARN", f" URL : {page.url}")
|
||
log("WARN", f" 标题 : {page.title}")
|
||
# 错误提示（timeout=0 不阻塞）
|
||
for sel in [".el-message--error", ".error-msg"]:
|
||
el = page.ele(sel, timeout=0)
|
||
if el and el.text.strip():
|
||
log("WARN", f" 页面提示: {el.text.strip()}")
|
||
log("WARN", "─────────────────────────────")
|
||
|
||
|
||
# ── 浏览器 ────────────────────────────────────────────────────────────────────
|
||
def get_page(headless: bool = False, port: int = 9222) -> ChromiumPage:
|
||
co = ChromiumOptions()
|
||
# 强制在生产环境下使用无头模式
|
||
is_docker = is_docker_env()
|
||
|
||
if headless or is_docker:
|
||
co.set_argument("--headless=new")
|
||
co.set_argument("--disable-gpu") # Docker 无头模式下强烈建议禁用 GPU，防止 404 和渲染崩溃
|
||
|
||
co.set_argument("--disable-blink-features=AutomationControlled")
|
||
co.set_argument("--no-sandbox")
|
||
co.set_argument("--disable-dev-shm-usage") # 防止 Docker 共享内存耗尽导致浏览器崩溃
|
||
co.set_argument("--disable-software-rasterizer") # 配合无头模式禁用软件光栅化器
|
||
co.set_argument("--remote-allow-origins=*") # 解决 Docker 下 websocket 404 问题
|
||
co.set_argument("--remote-debugging-address=127.0.0.1")
|
||
co.set_argument("--disable-web-security")
|
||
co.set_argument("--ignore-certificate-errors")
|
||
co.set_argument("--proxy-server=direct://") # 禁用代理
|
||
co.set_argument("--proxy-bypass-list=*")
|
||
co.set_argument("--window-size=1440,900")
|
||
|
||
if is_docker:
|
||
# Docker 生产环境：每次启动分配随机端口，避免前一个僵尸进程占用导致 404
|
||
co.auto_port(True)
|
||
# 为当前 worker 隔离 Chromium profile，防止多个进程共享 profile 导致握手错乱
|
||
user_data_dir = get_docker_user_data_dir()
|
||
co.set_user_data_path(str(user_data_dir))
|
||
log("INFO", f"Docker Chromium 用户目录: {user_data_dir}")
|
||
# 很多 Debian/Ubuntu 系统的 Chromium 实际上是通过 wrapper 脚本调用的
|
||
# 直接指定确切的执行路径，防止 DrissionPage 底层启动失败
|
||
if os.path.exists('/usr/bin/chromium'):
|
||
co.set_browser_path('/usr/bin/chromium')
|
||
elif os.path.exists('/usr/bin/chromium-browser'):
|
||
co.set_browser_path('/usr/bin/chromium-browser')
|
||
elif os.path.exists('/usr/bin/google-chrome'):
|
||
co.set_browser_path('/usr/bin/google-chrome')
|
||
else:
|
||
# 本地开发环境：使用固定端口，方便复用
|
||
co.set_local_port(port)
|
||
|
||
try:
|
||
page = ChromiumPage(co)
|
||
return page
|
||
except Exception as e:
|
||
log("ERR", f"浏览器初始化失败: {e}")
|
||
raise
|
||
|
||
|
||
# ── Vue 表单专用输入（JS setter + 模拟键盘） ──────────────────────────────────
|
||
def set_input_value(page: ChromiumPage, ele, value: str):
|
||
"""
|
||
Vue 双向绑定下 clear()+input() 会留残值，必须用 JS 原生 setter 清空
|
||
再通过模拟键盘输入触发 keydown/change 事件。
|
||
"""
|
||
page.run_js("""
|
||
var el = arguments[0];
|
||
Object.getOwnPropertyDescriptor(
|
||
window.HTMLInputElement.prototype, 'value'
|
||
).set.call(el, '');
|
||
el.dispatchEvent(new Event('input', {bubbles: true}));
|
||
""", ele)
|
||
time.sleep(0.05) # 等 Vue 处理 clear 事件
|
||
ele.input(value)
|
||
|
||
|
||
# ── 自动登录（使用 .env 账号） ────────────────────────────────────────────────
|
||
def login(page: ChromiumPage) -> bool:
|
||
"""自动填写 .env 中的账号信息并登录。"""
|
||
log("INFO", f"打开登录页: {ERP_URL}")
|
||
page.get(ERP_URL)
|
||
|
||
# 先用较短 timeout 等登录表单；若不出现则检测是否已处于登录状态
|
||
tenant_input = page.ele("@placeholder=租户代码", timeout=5)
|
||
if not tenant_input:
|
||
# 没有登录表单 → 判断是否已登录（主界面无密码框）
|
||
has_pwd = bool(page.ele("@placeholder=密码", timeout=0))
|
||
if not has_pwd:
|
||
log("OK", f"检测到已登录状态，跳过登录 → {page.url}")
|
||
return True
|
||
log("ERR", "未找到登录表单，页面可能未加载")
|
||
dump_page_state(page, "找不到表单")
|
||
return False
|
||
|
||
log("INFO", "填写登录信息...")
|
||
set_input_value(page, tenant_input, ERP_TENANT)
|
||
set_input_value(page, page.ele("@placeholder=账号"), ERP_USERNAME)
|
||
set_input_value(page, page.ele("@placeholder=密码"), ERP_PASSWORD)
|
||
time.sleep(0.2) # 等所有字段的 change/blur 事件结算
|
||
|
||
# 点击登录按钮
|
||
login_btn = page.ele(
|
||
"xpath://*[@id='app']/div/div[3]/div/div[2]/div[2]/div[5]/button",
|
||
timeout=5
|
||
)
|
||
if not login_btn:
|
||
log("ERR", "未找到登录按钮（XPath 失效，请更新选择器）")
|
||
dump_page_state(page, "找不到登录按钮")
|
||
return False
|
||
|
||
log("INFO", "点击登录...")
|
||
login_btn.click()
|
||
|
||
return _wait_for_login(page)
|
||
|
||
|
||
# ── 手动登录（等待用户在浏览器操作） ─────────────────────────────────────────
|
||
def login_manual(page: ChromiumPage) -> bool:
|
||
"""打开登录页，等待用户手动完成登录操作。"""
|
||
log("INFO", f"打开登录页（手动模式）: {ERP_URL}")
|
||
page.get(ERP_URL)
|
||
|
||
# 确认登录页已加载
|
||
tenant_input = page.ele("@placeholder=租户代码", timeout=15)
|
||
if not tenant_input:
|
||
log("ERR", "未找到登录表单")
|
||
dump_page_state(page, "找不到表单")
|
||
return False
|
||
|
||
log("INFO", "═" * 50)
|
||
log("INFO", " 请在浏览器中手动输入账号并登录")
|
||
log("INFO", " 程序将在检测到登录成功后自动继续...")
|
||
log("INFO", "═" * 50)
|
||
|
||
return _wait_for_login(page, timeout=120) # 手动模式等待 2 分钟
|
||
|
||
|
||
# ── 等待登录结果（自动/手动共用） ────────────────────────────────────────────
|
||
def _wait_for_login(page: ChromiumPage, timeout: int = 15) -> bool:
|
||
"""轮询检测登录结果：表单消失 → 成功；error 提示 → 失败。"""
|
||
for elapsed in range(1, timeout + 1):
|
||
time.sleep(1)
|
||
has_form = bool(page.ele("@placeholder=密码", timeout=0))
|
||
err_el = page.ele(".el-message--error", timeout=0)
|
||
err_text = err_el.text.strip() if err_el else ""
|
||
|
||
if err_text:
|
||
log("ERR", f"登录失败: {err_text}")
|
||
dump_page_state(page)
|
||
return False
|
||
|
||
if not has_form:
|
||
log("OK", f"登录成功（{elapsed}s）→ {page.url}")
|
||
return True
|
||
|
||
if elapsed % 10 == 0: # 每 10 秒提示一次进度（手动模式有用）
|
||
log("INFO", f" 等待登录中... ({elapsed}s)")
|
||
|
||
log("WARN", f"超过 {timeout}s 未检测到登录成功")
|
||
dump_page_state(page, "登录超时")
|
||
return False
|
||
|
||
|
||
# ── 单独运行时的入口 ──────────────────────────────────────────────────────────
|
||
if __name__ == "__main__":
|
||
import sys
|
||
manual = "--manual" in sys.argv
|
||
page = get_page()
|
||
try:
|
||
ok = login_manual(page) if manual else login(page)
|
||
if ok:
|
||
log("OK", "登录完成，按 Enter 关闭浏览器")
|
||
input()
|
||
else:
|
||
log("ERR", "登录失败")
|
||
time.sleep(3)
|
||
except KeyboardInterrupt:
|
||
log("INFO", "用户中断")
|
||
finally:
|
||
page.quit()
|
||
log("INFO", "浏览器已关闭")
|