udxandCursor
9793d8581d
Bring in plaza props/prefabs, ZhaoYun source GLB, design originals, and Hunyuan/Blender tooling used to produce them. Co-authored-by: Cursor <cursoragent@cursor.com>
336 lines
11 KiB
Python
336 lines
11 KiB
Python
"""WeChat login session for https://3d.hunyuan.tencent.com via Playwright.
|
||
|
||
CLI 友好：同步阻塞扫码登录（非后台线程）。
|
||
未登录 / 凭证失效 / 额度满换号 时弹出有头 Edge 供用户扫码。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Any, Callable, Optional
|
||
|
||
HUNYUAN_URL = "https://3d.hunyuan.tencent.com/"
|
||
LOGIN_TIMEOUT_SEC = 300
|
||
|
||
UpdateFn = Callable[..., None]
|
||
|
||
|
||
def _tools_root() -> Path:
|
||
return Path(__file__).resolve().parent.parent
|
||
|
||
|
||
def session_dir() -> Path:
|
||
override = os.environ.get("HUNYUAN_SESSION_DIR", "").strip()
|
||
if override:
|
||
return Path(override).expanduser().resolve()
|
||
return _tools_root() / "data" / "hunyuan"
|
||
|
||
|
||
def storage_state_path() -> Path:
|
||
return session_dir() / "storage_state.json"
|
||
|
||
|
||
def status_file_path() -> Path:
|
||
return session_dir() / "status.json"
|
||
|
||
|
||
def jobs_dir() -> Path:
|
||
path = session_dir() / "jobs"
|
||
path.mkdir(parents=True, exist_ok=True)
|
||
return path
|
||
|
||
|
||
def _browser_channel() -> str:
|
||
return os.environ.get("HUNYUAN_BROWSER_CHANNEL", "msedge").strip() or "msedge"
|
||
|
||
|
||
def _write_status(**values: Any) -> dict:
|
||
directory = session_dir()
|
||
directory.mkdir(parents=True, exist_ok=True)
|
||
payload = status()
|
||
payload.update(values, updated_at=time.time())
|
||
temporary = status_file_path().with_suffix(".tmp")
|
||
temporary.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
||
temporary.replace(status_file_path())
|
||
return payload
|
||
|
||
|
||
def status() -> dict:
|
||
path = status_file_path()
|
||
if path.is_file():
|
||
try:
|
||
data = json.loads(path.read_text("utf-8"))
|
||
if storage_state_path().is_file() and data.get("state") not in {
|
||
"waiting_scan",
|
||
"error",
|
||
"timeout",
|
||
}:
|
||
data["logged_in"] = True
|
||
return data
|
||
except (ValueError, OSError):
|
||
pass
|
||
logged = storage_state_path().is_file()
|
||
return {
|
||
"logged_in": logged,
|
||
"state": "ready" if logged else "logged_out",
|
||
"message": "已登录混元3D" if logged else "尚未微信登录",
|
||
"session_dir": str(session_dir()),
|
||
}
|
||
|
||
|
||
def is_logged_in() -> bool:
|
||
path = storage_state_path()
|
||
if not path.is_file() or path.stat().st_size < 8:
|
||
return False
|
||
try:
|
||
with path.open("rb") as fp:
|
||
if fp.read(1) != b"{":
|
||
return False
|
||
json.loads(path.read_text("utf-8"))
|
||
return True
|
||
except (OSError, ValueError, UnicodeDecodeError):
|
||
return False
|
||
|
||
|
||
def logout() -> dict:
|
||
for path in (storage_state_path(), status_file_path()):
|
||
if path.is_file():
|
||
path.unlink()
|
||
return _write_status(logged_in=False, state="logged_out", message="已退出混元3D 登录")
|
||
|
||
|
||
def _click_first(page, labels, exact=False) -> bool:
|
||
for label in labels:
|
||
candidates = [
|
||
page.get_by_role("button", name=label),
|
||
page.get_by_role("link", name=label),
|
||
page.get_by_text(label, exact=exact),
|
||
]
|
||
for locator in candidates:
|
||
try:
|
||
if locator.count() and locator.first.is_visible(timeout=800):
|
||
locator.first.click(timeout=3000)
|
||
page.wait_for_timeout(600)
|
||
return True
|
||
except Exception:
|
||
continue
|
||
return False
|
||
|
||
|
||
def _looks_logged_in(page) -> bool:
|
||
try:
|
||
if page.get_by_text("登录", exact=True).count():
|
||
login = page.get_by_text("登录", exact=True).first
|
||
if login.is_visible(timeout=300):
|
||
return False
|
||
except Exception:
|
||
pass
|
||
cookies = page.context.cookies()
|
||
interesting = [
|
||
cookie
|
||
for cookie in cookies
|
||
if any(
|
||
token in cookie["name"].lower()
|
||
for token in ("token", "openid", "uin", "skey", "hy3d", "session", "userid")
|
||
)
|
||
]
|
||
try:
|
||
storage_keys = page.evaluate(
|
||
"() => Object.keys({...localStorage, ...sessionStorage})"
|
||
)
|
||
except Exception:
|
||
storage_keys = []
|
||
storage_hit = any(
|
||
any(token in key.lower() for token in ("token", "user", "auth", "openid", "uin"))
|
||
for key in storage_keys
|
||
)
|
||
has_generate = False
|
||
try:
|
||
has_generate = bool(page.get_by_text("立即生成").count())
|
||
except Exception:
|
||
pass
|
||
# 有生成入口，且 cookie/storage 任一命中，或页面上已无「登录」
|
||
if has_generate and (interesting or storage_hit):
|
||
return True
|
||
if has_generate and not _page_shows_login_cta(page):
|
||
return True
|
||
return bool(interesting and storage_hit)
|
||
|
||
|
||
def _page_shows_login_cta(page) -> bool:
|
||
try:
|
||
if page.get_by_text("登录", exact=True).count():
|
||
return page.get_by_text("登录", exact=True).first.is_visible(timeout=300)
|
||
except Exception:
|
||
pass
|
||
return False
|
||
|
||
|
||
def _try_site_logout(page) -> None:
|
||
"""尽量在官网侧退出，方便额度满后换号扫码。"""
|
||
for label in ("退出登录", "退出", "登出", "注销"):
|
||
if _click_first(page, (label,), exact=False):
|
||
page.wait_for_timeout(800)
|
||
break
|
||
# 常见头像/用户菜单再点一次退出
|
||
try:
|
||
avatar = page.locator(
|
||
'[class*="avatar"], [class*="user-info"], [class*="UserInfo"]'
|
||
).first
|
||
if avatar.count() and avatar.is_visible(timeout=500):
|
||
avatar.click(timeout=2000)
|
||
page.wait_for_timeout(400)
|
||
_click_first(page, ("退出登录", "退出", "登出"), exact=False)
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def start_wechat_login_blocking(
|
||
timeout_sec: int = LOGIN_TIMEOUT_SEC,
|
||
*,
|
||
update: Optional[UpdateFn] = None,
|
||
reason: str = "auth",
|
||
clear_existing: bool = False,
|
||
) -> dict:
|
||
"""打开有头浏览器供微信扫码，阻塞直到成功 / 超时 / 失败。
|
||
|
||
reason:
|
||
- auth: 未登录或凭证失效
|
||
- quota: 当前账号额度满，需换号
|
||
clear_existing: 先清本地会话；quota 场景会再尝试官网退出。
|
||
"""
|
||
from playwright.sync_api import sync_playwright
|
||
|
||
def _emit(event: str, **payload: Any) -> None:
|
||
if update:
|
||
update(event=event, **payload)
|
||
|
||
if clear_existing or reason == "quota":
|
||
logout()
|
||
|
||
if reason == "quota":
|
||
message = (
|
||
"当前混元账号额度/次数已满，请在弹出的浏览器中退出旧号并用微信扫码登录其他账号"
|
||
)
|
||
else:
|
||
message = "请在弹出的浏览器窗口中用微信扫码登录混元3D"
|
||
|
||
_write_status(
|
||
logged_in=False,
|
||
state="waiting_scan",
|
||
message=message,
|
||
reason=reason,
|
||
)
|
||
_emit(
|
||
"login_waiting",
|
||
stage=message,
|
||
progress=5,
|
||
timeout_sec=timeout_sec,
|
||
reason=reason,
|
||
)
|
||
# 同步打到 stderr，方便 Agent/终端注意弹窗
|
||
print(f"[hunyuan] {message}（最长等待 {timeout_sec}s）", flush=True)
|
||
|
||
try:
|
||
with sync_playwright() as playwright:
|
||
browser = playwright.chromium.launch(
|
||
channel=_browser_channel(), headless=False
|
||
)
|
||
context = browser.new_context(
|
||
locale="zh-CN", viewport={"width": 1280, "height": 900}
|
||
)
|
||
page = context.new_page()
|
||
page.goto(HUNYUAN_URL, wait_until="domcontentloaded", timeout=60000)
|
||
page.wait_for_timeout(2000)
|
||
|
||
if reason == "quota":
|
||
_try_site_logout(page)
|
||
page.goto(HUNYUAN_URL, wait_until="domcontentloaded", timeout=60000)
|
||
page.wait_for_timeout(1500)
|
||
|
||
_click_first(page, ("登录", "立即登录"), exact=True)
|
||
page.wait_for_timeout(800)
|
||
_click_first(page, ("微信", "微信登录", "微信扫码"), exact=False)
|
||
|
||
deadline = time.monotonic() + timeout_sec
|
||
while time.monotonic() < deadline:
|
||
if _looks_logged_in(page):
|
||
session_dir().mkdir(parents=True, exist_ok=True)
|
||
context.storage_state(path=str(storage_state_path()))
|
||
result = _write_status(
|
||
logged_in=True,
|
||
state="ready",
|
||
message="微信登录成功",
|
||
reason=reason,
|
||
)
|
||
browser.close()
|
||
_emit("login_ok", stage="微信登录成功", progress=15, **result)
|
||
print("[hunyuan] 微信登录成功，会话已保存", flush=True)
|
||
return result
|
||
page.wait_for_timeout(1500)
|
||
|
||
result = _write_status(
|
||
logged_in=False,
|
||
state="timeout",
|
||
message="微信登录超时，请重试",
|
||
reason=reason,
|
||
)
|
||
browser.close()
|
||
_emit("login_failed", stage="微信登录超时", progress=0, **result)
|
||
raise TimeoutError("微信登录超时，请重试")
|
||
except TimeoutError:
|
||
raise
|
||
except Exception as exc:
|
||
result = _write_status(
|
||
logged_in=False, state="error", message=f"登录失败：{exc}", reason=reason
|
||
)
|
||
_emit("login_failed", stage=str(exc), progress=0, **result)
|
||
raise RuntimeError(f"登录失败：{exc}") from exc
|
||
|
||
|
||
def ensure_login(
|
||
*,
|
||
interactive: bool = True,
|
||
timeout_sec: int = LOGIN_TIMEOUT_SEC,
|
||
update: Optional[UpdateFn] = None,
|
||
force: bool = False,
|
||
reason: str = "auth",
|
||
) -> dict:
|
||
"""若无会话则弹微信登录；interactive=False 且未登录则抛错。
|
||
|
||
force=True 时忽略本地凭证，强制打开登录页（换号 / 重登）。
|
||
"""
|
||
if is_logged_in() and not force:
|
||
current = status()
|
||
if update:
|
||
update(
|
||
event="login_cached",
|
||
stage="已有混元登录凭证",
|
||
progress=12,
|
||
**{k: v for k, v in current.items() if k != "event"},
|
||
)
|
||
return current
|
||
|
||
if not interactive:
|
||
raise RuntimeError("尚未微信登录混元3D，请先执行 --login 或允许交互登录。")
|
||
|
||
if update:
|
||
update(
|
||
event="login_required",
|
||
stage=(
|
||
"当前账号额度已满，需要换号登录"
|
||
if reason == "quota"
|
||
else "需要微信登录混元3D"
|
||
),
|
||
progress=2,
|
||
reason=reason,
|
||
)
|
||
return start_wechat_login_blocking(
|
||
timeout_sec=timeout_sec,
|
||
update=update,
|
||
reason=reason,
|
||
clear_existing=force or reason == "quota",
|
||
)
|