Spaces:
Running
Running
Download src/playwright_worker.py from jzyg123/grade_query: direct link, hf CLI and curl.
- Browser
- Download file 16.2 kB
-
https://huggingface.co/spaces/jzyg123/grade_query/resolve/main/src/playwright_worker.py
- Command line
-
hf download hf://spaces/jzyg123/grade_query/src/playwright_worker.py
-
curl -L -o playwright_worker.py https://huggingface.co/spaces/jzyg123/grade_query/resolve/main/src/playwright_worker.py
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 | |