grade_query / src /playwright_worker.py
jzyg123's picture
Upload 4 files
35939d5 verified
Raw History Blame Contribute Delete
16.2 kB
from flask import Flask, jsonify, render_template, request, send_from_directory
from playwright.sync_api import sync_playwright, TimeoutError as PlaywrightTimeoutError
from pathlib import Path
import base64
import json
import os
import subprocess
import threading
import time
BASE_DIR = Path.cwd().resolve()
STATIC_DIR = BASE_DIR / "static"
TEMPLATE_DIR = BASE_DIR / "templates"
DATA_DIR = BASE_DIR / "data"
STATIC_DIR.mkdir(parents=True, exist_ok=True)
TEMPLATE_DIR.mkdir(parents=True, exist_ok=True)
DATA_DIR.mkdir(parents=True, exist_ok=True)
TARGET_PAGE = "https://ynsdfz.aliwork.com/APP_DKQG7TSUK6ZNU3RDFGVO/workbench"
TARGET_REQUEST_KEYWORD = "/dingtalk/web/APP_DKQG7TSUK6ZNU3RDFGVO/v1/form/searchFormDatas.json"
DINGTALK_CHALLENGE_URL_KEYWORD = "https://login.dingtalk.com/oauth2/challenge.htm"
DINGTALK_LOGIN_HOST_KEYWORD = "login.dingtalk.com"
CORP_TEXT = "云南师范大学附属中学"
QUERY_VALUE = "53010320080506463X"
LOGIN_WAIT_SECONDS = 120
COOKIE_FILE = BASE_DIR / "cookies.txt"
HEADERS_FILE = DATA_DIR / "request_headers_debug.json"
STORAGE_STATE_FILE = DATA_DIR / "playwright_state.json"
QR_IMAGE_FILE = STATIC_DIR / "dingtalk_qr.png"
SCREENSHOT_FILE = STATIC_DIR / "browser_screenshot.png"
SHOW_COOKIE_IN_FRONTEND = False
state_lock = threading.Lock()
worker_lock = threading.Lock()
worker_thread = None
worker_runtime = {
"playwright": None,
"browser": None,
"context": None,
"page": None,
}
STATE = {
"running": False,
"status": "idle",
"current_url": "",
"cookie": "",
"cookie_source": "",
"request_url": "",
"error": "",
"last_updated": 0,
"logs": [],
"has_qr": False,
"has_screenshot": False,
}
def set_state(**kwargs):
with state_lock:
STATE.update(kwargs)
STATE["last_updated"] = int(time.time())
def add_log(message, level="info"):
import datetime
ts = datetime.datetime.now().strftime("%H:%M:%S.%f")[:-3]
with state_lock:
STATE["logs"].append({"time": ts, "message": message, "level": level})
STATE["logs"] = STATE["logs"][-300:]
STATE["last_updated"] = int(time.time())
def get_state_with_options(include_sensitive=True):
with state_lock:
data = dict(STATE)
data["logs"] = list(STATE["logs"])
data["has_qr"] = QR_IMAGE_FILE.exists()
data["has_screenshot"] = SCREENSHOT_FILE.exists()
data["show_cookie_in_frontend"] = SHOW_COOKIE_IN_FRONTEND
if not include_sensitive:
data["cookie"] = ""
return data
def get_state(include_sensitive=True):
return get_state_with_options(include_sensitive=include_sensitive)
def build_cookie_header(cookies):
parts = []
for c in cookies:
name = c.get("name")
value = c.get("value")
if name and value is not None:
parts.append(f"{name}={value}")
return "; ".join(parts)
def save_text(path: Path, content: str):
temp_path = path.with_name(f"{path.name}.tmp")
temp_path.write_text(content, encoding="utf-8")
temp_path.replace(path)
def save_json(path: Path, data):
save_text(path, json.dumps(data, ensure_ascii=False, indent=2))
def reset_runtime_artifacts():
for path in (QR_IMAGE_FILE, SCREENSHOT_FILE):
try:
if path.exists():
path.unlink()
except Exception:
pass
def click_corp_if_present(page):
try:
if DINGTALK_CHALLENGE_URL_KEYWORD not in page.url:
return False
corp_item = page.locator("div.module-corp-sel-listitem", has_text=CORP_TEXT).first
if corp_item.count() > 0 and corp_item.is_visible(timeout=1000):
add_log(f"检测到企业选择项:{CORP_TEXT},正在点击")
corp_item.click()
return True
except Exception:
pass
return False
def ensure_auto_login_checked(page):
try:
if DINGTALK_LOGIN_HOST_KEYWORD not in page.url:
return False
op_item = page.locator("div.module-qrcode-op-item", has_text="自动登录").first
if op_item.count() == 0 or not op_item.is_visible(timeout=1000):
return False
icon = op_item.locator("div.base-comp-check-box-rememberme-box").first
icon_class = icon.get_attribute("class") or ""
if "checkbox-undone" in icon_class:
add_log("检测到“自动登录”未勾选,正在勾选")
op_item.click()
time.sleep(0.3)
return True
except Exception:
pass
return False
def export_login_qr_if_present(page):
try:
if DINGTALK_LOGIN_HOST_KEYWORD not in page.url:
return False
canvas = page.locator("div.module-qrcode-code canvas").first
if canvas.count() == 0 or not canvas.is_visible(timeout=1000):
return False
data_url = canvas.evaluate("el => el.toDataURL('image/png')")
if not data_url or not data_url.startswith("data:image/png;base64,"):
return False
png_base64 = data_url.split(",", 1)[1]
png_bytes = base64.b64decode(png_base64)
if len(png_bytes) < 500:
add_log("检测到二维码图片可能为空白或未明确加载", "warning")
return "blank"
previous = QR_IMAGE_FILE.read_bytes() if QR_IMAGE_FILE.exists() else None
if previous == png_bytes:
return True
QR_IMAGE_FILE.write_bytes(png_bytes)
set_state(has_qr=True)
add_log(f"已提取登录二维码:{QR_IMAGE_FILE.name}", "success")
return True
except Exception as e:
add_log(f"提取登录二维码异常:{e}", "error")
return False
def click_immediate_login_if_present(page):
try:
# 检测是否出现快速登录提示
notice = page.get_by_text("此账号已在使用,可直接登录宜搭").first
if notice.count() > 0 and notice.is_visible(timeout=500):
# 找到立即登录按钮
btn = page.get_by_role("button", name="立即登录").first
if btn.count() == 0:
# 备用定位方式
btn = page.locator("text=立即登录").first
if btn.count() > 0 and btn.is_visible(timeout=500):
add_log("检测到快速登录提示,正在点击“立即登录”")
btn.click()
time.sleep(1)
return True
except Exception as e:
pass
return False
def is_query_page_ready(page):
try:
input_box = page.get_by_placeholder("请输入").first
if input_box.count() > 0 and input_box.is_visible(timeout=500):
query_button = page.get_by_role("button", name="查 询")
return query_button.count() > 0 and query_button.first.is_visible(timeout=500)
except Exception:
return False
return False
def click_qr_login_tab_if_present(page):
try:
if DINGTALK_LOGIN_HOST_KEYWORD not in page.url:
return False
tab = page.locator("div.dtm-segment-item", has_text="扫码登录").first
if tab.count() > 0 and tab.is_visible(timeout=500):
if tab.get_attribute("aria-selected") != "true":
add_log("检测到未处于扫码登录状态,正在切换")
tab.click()
time.sleep(0.5)
return True
except Exception:
pass
return False
def wait_until_query_page_ready(page, start_time=None, timeout_seconds=300):
if start_time is None:
start_time = time.time()
deadline = start_time + timeout_seconds
last_url = None
while time.time() < deadline:
try:
current_url = page.url
except Exception:
current_url = ""
set_state(current_url=current_url)
if current_url != last_url:
add_log(f"当前页面:{current_url}")
last_url = current_url
click_immediate_login_if_present(page)
click_corp_if_present(page)
click_qr_login_tab_if_present(page)
ensure_auto_login_checked(page)
qr_status = export_login_qr_if_present(page)
if qr_status == "blank":
return "blank_qr"
if is_query_page_ready(page):
return "ready"
try:
page.wait_for_load_state("domcontentloaded", timeout=1000)
except Exception:
pass
time.sleep(1)
return "timeout"
def cleanup_runtime():
runtime = worker_runtime
for key in ["page", "context", "browser", "playwright"]:
obj = runtime.get(key)
if obj is None:
continue
try:
if key == "playwright":
obj.stop()
else:
obj.close()
except Exception:
pass
runtime[key] = None
def automation_worker():
start_time = time.time()
with worker_lock:
max_attempts = 2
for attempt in range(max_attempts):
if time.time() - start_time > 300:
add_log("任务总耗时超过5分钟,自动终止此轮任务", "error")
set_state(status="timeout", error="任务总耗时超过5分钟", running=False)
break
if attempt == 0:
reset_runtime_artifacts()
set_state(
running=True,
status="starting",
error="",
current_url="",
cookie="",
cookie_source="",
request_url=""
)
add_log(f"自动化任务启动(第 {attempt + 1} 次尝试)")
success = False
try:
pw = sync_playwright().start()
worker_runtime["playwright"] = pw
browser = pw.chromium.launch(headless=True)
worker_runtime["browser"] = browser
if STORAGE_STATE_FILE.exists():
if attempt == 0:
context = browser.new_context(storage_state=str(STORAGE_STATE_FILE))
add_log(f"已加载登录状态:{STORAGE_STATE_FILE.name}")
else:
try:
STORAGE_STATE_FILE.unlink()
except Exception:
pass
add_log("已尝试清空旧的登录状态数据")
context = browser.new_context()
add_log("使用新会话(未加载已有状态)")
else:
context = browser.new_context()
add_log("未发现已保存登录状态,使用新会话")
worker_runtime["context"] = context
page = context.new_page()
worker_runtime["page"] = page
set_state(status="opening")
add_log("正在打开目标页面")
page.goto(TARGET_PAGE, wait_until="load")
set_state(status="waiting_login")
add_log("正在自动检测登录完成状态")
time_elapsed = time.time() - start_time
remaining_time = max(10, 300 - time_elapsed)
ready_status = wait_until_query_page_ready(page, start_time=time.time(), timeout_seconds=remaining_time)
if ready_status == "timeout":
add_log("等待登录完成超时", "warning" if attempt == 0 else "error")
if attempt == 0:
continue
else:
set_state(status="login_timeout", error="在限定时间内未进入查询页面")
break
elif ready_status == "blank_qr":
add_log("二维码持续空白或异常", "warning" if attempt == 0 else "error")
if attempt == 0:
continue
else:
set_state(status="login_timeout", error="二维码为空白")
break
elif not ready_status or ready_status == False:
if attempt == 0:
continue
else:
set_state(status="login_timeout", error="未能进入查询页面")
break
context.storage_state(path=str(STORAGE_STATE_FILE))
add_log(f"已保存登录状态:{STORAGE_STATE_FILE.name}", "success")
set_state(status="querying")
input_box = page.get_by_placeholder("请输入").first
input_box.click()
input_box.fill(QUERY_VALUE)
add_log(f"已填入查询值:{QUERY_VALUE}")
add_log("准备点击查询,并同步等待目标请求")
with page.expect_request(lambda req: TARGET_REQUEST_KEYWORD in req.url, timeout=30000) as req_info:
query_button = page.get_by_role("button", name="查 询")
query_button.click()
request_obj = req_info.value
request_url = request_obj.url
headers = {}
try:
headers = request_obj.all_headers()
except Exception:
try:
headers = request_obj.headers
except Exception:
headers = {}
debug_headers = dict(headers or {})
if "cookie" in debug_headers:
debug_headers["cookie"] = "[saved-to-cookies.txt]"
save_json(
HEADERS_FILE,
{
"url": request_url,
"headers": debug_headers,
"cookie_header_captured": bool((headers or {}).get("cookie"))
}
)
cookie = headers.get("cookie") if headers else None
source = "直接从目标请求头读取"
if not cookie:
try:
cookies = context.cookies([request_url])
except Exception:
cookies = context.cookies()
cookie = build_cookie_header(cookies)
source = "按目标请求可发送规则重建"
if cookie:
save_text(COOKIE_FILE, cookie)
add_log(f"已覆盖保存本地 Cookie 文件:{COOKIE_FILE.name}", "success")
set_state(cookie=cookie, cookie_source=source, request_url=request_url)
add_log(f"已成功获取 cookie,来源:{source}", "success")
set_state(status="done")
else:
set_state(status="cookie_missing", request_url=request_url, error="未能拿到 cookie")
add_log("已捕获目标请求,但未能拿到 cookie", "error")
success = True
break # Successful, exit the loop
except PlaywrightTimeoutError:
set_state(status="request_timeout", error="30 秒内未捕获到目标请求")
add_log("等待目标请求超时", "error")
if attempt == 0:
add_log("准备重试...", "warning")
except Exception as e:
set_state(status="error", error=str(e))
add_log(f"发生异常:{e}", "error")
if attempt == 0:
add_log("准备重试...", "warning")
finally:
cleanup_runtime()
if success or attempt == max_attempts - 1:
set_state(running=False)
add_log("自动化任务结束")
def start_worker_if_needed():
global worker_thread
if worker_thread and worker_thread.is_alive():
return False
worker_thread = threading.Thread(target=automation_worker, daemon=True)
worker_thread.start()
return True