feat(rpa_helpers): add auth_flows.py with 9 ensure_logged_in_* and dispatcher

Co-authored-by: Cursor <cursoragent@cursor.com>

scripts/service/rpa_helpers/auth_flows.py

@@ -0,0 +1,449 @@
+# -*- coding: utf-8 -*-
+"""Auth strategy helpers: ensure_logged_in_* + dispatcher (no DB / no CLI)."""
+from __future__ import annotations
+
+import hashlib
+import json
+import re
+import socket
+import sys
+import time
+import uuid
+from dataclasses import dataclass
+from urllib.parse import urlparse
+
+from .human import human_click, human_type
+
+
+def _poll_until(
+    predicate,
+    timeout_sec: int,
+    interval_sec: float = 2.0,
+) -> bool:
+    """轮询 predicate 直到返回 True 或超时。"""
+    deadline = time.time() + timeout_sec
+    while time.time() < deadline:
+        try:
+            if predicate():
+                return True
+        except Exception:
+            pass
+        time.sleep(interval_sec)
+    return False
+
+
+def _stderr_prompt(msg: str) -> None:
+    sys.stderr.write(f"[auth_flows] {msg}\n")
+    sys.stderr.flush()
+
+
+@dataclass
+class LoginResult:
+    ok: bool
+    needs_human: bool
+    message: str
+
+
+class AuthStrategyError(Exception):
+    """auth_strategy 相关错误（不可自动化 / 不支持 / fingerprint 不匹配等）。"""
+
+    def __init__(self, code: str, message: str):
+        super().__init__(message)
+        self.code = code
+        self.message = message
+
+
+def _account_extra(account: dict) -> dict:
+    ex = account.get("extra_json")
+    if ex is None:
+        return {}
+    if isinstance(ex, dict):
+        return ex
+    if isinstance(ex, str):
+        try:
+            return json.loads(ex)
+        except json.JSONDecodeError:
+            return {}
+    return {}
+
+
+def _is_already_logged_in(page, account: dict) -> bool:
+    """粗略判断是否已登录（可被 extra_json.logged_in_url_pattern 覆盖）。"""
+    try:
+        cur = page.url or ""
+    except Exception:
+        return False
+    extra = _account_extra(account)
+    pattern = extra.get("logged_in_url_pattern")
+    if isinstance(pattern, str) and pattern:
+        try:
+            return bool(re.search(pattern, cur))
+        except re.error:
+            pass
+    lc = cur.lower()
+    if any(k in lc for k in ("login", "signin", "sign-in", "auth", "/sso/")):
+        return False
+    try:
+        login_btn = page.locator(
+            "button:has-text('登录'), button:has-text('Sign in'), "
+            "a:has-text('登录'), a:has-text('Sign in')"
+        )
+        if login_btn.count() > 0:
+            try:
+                if login_btn.first.is_visible(timeout=500):
+                    return False
+            except Exception:
+                pass
+    except Exception:
+        pass
+    return True
+
+
+def _current_machine_fingerprint() -> str:
+    """hostname + MAC 的 sha256 前 16 位。"""
+    h = hashlib.sha256()
+    h.update(socket.gethostname().encode("utf-8"))
+    h.update(str(uuid.getnode()).encode("utf-8"))
+    return h.hexdigest()[:16]
+
+
+def _login_selectors(account: dict) -> tuple[str, str, str]:
+    extra = _account_extra(account)
+    selectors = extra.get("login_selectors") or {}
+    login_id_sel = (
+        selectors.get("login_id")
+        or "input[name*='user'], input[name*='login'], input[name*='account'], input[type='email']"
+    )
+    password_sel = selectors.get("password") or "input[type='password']"
+    submit_sel = (
+        selectors.get("submit")
+        or "button[type='submit'], button:has-text('登录'), button:has-text('Sign in')"
+    )
+    return str(login_id_sel), str(password_sel), str(submit_sel)
+
+
+def _require_login_id(account: dict) -> str | None:
+    lid = account.get("login_id")
+    if lid is None:
+        return None
+    s = str(lid).strip()
+    return s or None
+
+
+def ensure_logged_in_password_auto(page, account: dict, password: str, **_kwargs) -> LoginResult:
+    if _is_already_logged_in(page, account):
+        return LoginResult(
+            ok=True,
+            needs_human=False,
+            message="already logged in (profile reused)",
+        )
+    url = account.get("url") or ""
+    if url:
+        page.goto(url, wait_until="domcontentloaded", timeout=60000)
+
+    login_id = _require_login_id(account)
+    if not login_id:
+        return LoginResult(ok=False, needs_human=False, message="login_id missing in account")
+
+    lid_sel, pwd_sel, sub_sel = _login_selectors(account)
+    human_type(page.locator(lid_sel).first, login_id)
+    human_type(page.locator(pwd_sel).first, password)
+    human_click(page.locator(sub_sel).first)
+
+    ok = _poll_until(lambda: _is_already_logged_in(page, account), 30, interval_sec=1.0)
+    return LoginResult(
+        ok=ok,
+        needs_human=False,
+        message="login succeeded" if ok else "login timeout",
+    )
+
+
+def ensure_logged_in_password_with_captcha(
+    page,
+    account: dict,
+    password: str,
+    *,
+    captcha_value: str | None = None,
+    wait_human_sec: int = 60,
+    **_kwargs,
+) -> LoginResult:
+    if _is_already_logged_in(page, account):
+        return LoginResult(ok=True, needs_human=False, message="already logged in (profile reused)")
+    url = account.get("url") or ""
+    if url:
+        page.goto(url, wait_until="domcontentloaded", timeout=60000)
+
+    login_id = _require_login_id(account)
+    if not login_id:
+        return LoginResult(ok=False, needs_human=False, message="login_id missing in account")
+
+    extra = _account_extra(account)
+    selectors = extra.get("login_selectors") or {}
+    captcha_sel = (
+        selectors.get("captcha") or "input[name*='captcha'], input[name*='verify_code']"
+    )
+    lid_sel, pwd_sel, sub_sel = _login_selectors(account)
+
+    human_type(page.locator(lid_sel).first, login_id)
+    human_type(page.locator(pwd_sel).first, password)
+
+    if captcha_value is not None:
+        human_type(page.locator(captcha_sel).first, str(captcha_value))
+        human_click(page.locator(sub_sel).first)
+        ok = _poll_until(lambda: _is_already_logged_in(page, account), wait_human_sec, interval_sec=1.0)
+        return LoginResult(
+            ok=ok,
+            needs_human=False,
+            message="login succeeded" if ok else "captcha flow timeout",
+        )
+
+    _stderr_prompt("请输入图形验证码后回车继续（CTRL+C 取消）")
+    ok = _poll_until(lambda: _is_already_logged_in(page, account), wait_human_sec, interval_sec=1.0)
+    return LoginResult(
+        ok=ok,
+        needs_human=True,
+        message="captcha human assist completed" if ok else "captcha human assist timeout",
+    )
+
+
+def _is_2fa_page(page) -> bool:
+    try:
+        u = (page.url or "").lower()
+        if any(x in u for x in ("2fa", "verify", "mfa")):
+            return True
+        otp = page.locator("input[name*='code'], input[name*='otp']")
+        if otp.count() > 0:
+            return True
+    except Exception:
+        pass
+    return False
+
+
+def ensure_logged_in_password_plus_2fa(
+    page,
+    account: dict,
+    password: str,
+    *,
+    wait_human_sec: int = 120,
+    **_kwargs,
+) -> LoginResult:
+    if _is_already_logged_in(page, account):
+        return LoginResult(ok=True, needs_human=False, message="already logged in (profile reused)")
+    url = account.get("url") or ""
+    if url:
+        page.goto(url, wait_until="domcontentloaded", timeout=60000)
+
+    login_id = _require_login_id(account)
+    if not login_id:
+        return LoginResult(ok=False, needs_human=False, message="login_id missing in account")
+
+    lid_sel, pwd_sel, sub_sel = _login_selectors(account)
+    human_type(page.locator(lid_sel).first, login_id)
+    human_type(page.locator(pwd_sel).first, password)
+    human_click(page.locator(sub_sel).first)
+
+    _poll_until(lambda: _is_2fa_page(page) or _is_already_logged_in(page, account), 30, interval_sec=1.0)
+
+    if _is_already_logged_in(page, account):
+        return LoginResult(ok=True, needs_human=False, message="login succeeded without 2FA")
+
+    _stderr_prompt("请输入 2FA 验证码并提交")
+    ok = _poll_until(lambda: _is_already_logged_in(page, account), wait_human_sec, interval_sec=1.0)
+    msg = "2FA completed" if ok else "2FA timeout ERR_LOGIN_2FA_TIMEOUT"
+    return LoginResult(ok=ok, needs_human=True, message=msg)
+
+
+def _qr_locator(page):
+    return page.locator("canvas, img[src*='qr'], img[alt*='二维码']")
+
+
+def ensure_logged_in_qr_code_manual(page, account: dict, *, wait_human_sec: int = 300, **_kwargs) -> LoginResult:
+    if _is_already_logged_in(page, account):
+        return LoginResult(ok=True, needs_human=False, message="already logged in (profile reused)")
+    url = account.get("url") or ""
+    if url:
+        page.goto(url, wait_until="domcontentloaded", timeout=60000)
+
+    qr = _qr_locator(page)
+    try:
+        qr.first.wait_for(state="visible", timeout=30000)
+    except Exception:
+        return LoginResult(ok=False, needs_human=True, message="qr code not found")
+
+    _stderr_prompt("请用手机扫描浏览器中的二维码")
+
+    def _qr_done() -> bool:
+        if _is_already_logged_in(page, account):
+            return True
+        try:
+            av = page.locator("img[alt*='avatar']")
+            if av.count() > 0 and av.first.is_visible(timeout=300):
+                return True
+        except Exception:
+            pass
+        return False
+
+    ok = _poll_until(_qr_done, wait_human_sec, interval_sec=2.0)
+    return LoginResult(
+        ok=ok,
+        needs_human=True,
+        message="qr scan completed" if ok else "qr scan timeout",
+    )
+
+
+def ensure_logged_in_per_session_manual(
+    page,
+    account: dict,
+    *,
+    wait_human_sec: int = 600,
+    **_kwargs,
+) -> LoginResult:
+    url = account.get("url") or ""
+    if url:
+        page.goto(url, wait_until="domcontentloaded", timeout=60000)
+
+    _stderr_prompt(f"请完整登录后等待 RPA 继续（最长 {wait_human_sec} 秒）")
+
+    ok = _poll_until(lambda: _is_already_logged_in(page, account), wait_human_sec, interval_sec=2.0)
+    return LoginResult(
+        ok=ok,
+        needs_human=True,
+        message="per-session manual completed" if ok else "per-session manual timeout",
+    )
+
+
+def ensure_logged_in_device_bound(
+    page,
+    account: dict,
+    password: str | None = None,
+    *,
+    wait_human_sec: int = 120,
+    **_kwargs,
+) -> LoginResult:
+    expected_fp = account.get("device_fingerprint")
+    if expected_fp:
+        cur = _current_machine_fingerprint()
+        if cur != str(expected_fp):
+            raise AuthStrategyError(
+                "ERR_DEVICE_FINGERPRINT_MISMATCH",
+                f"current machine fingerprint {cur!r} does not match account's {expected_fp!r}",
+            )
+    if password:
+        return ensure_logged_in_password_auto(page, account, password)
+    return ensure_logged_in_per_session_manual(page, account, wait_human_sec=wait_human_sec)
+
+
+def ensure_logged_in_client_certificate(page, account: dict, **_kwargs) -> LoginResult:
+    raise AuthStrategyError(
+        "ERR_AUTH_STRATEGY_NOT_AUTOMATABLE",
+        "client_certificate strategy requires manual cert + PIN; not automatable by rpa_helpers",
+    )
+
+
+def ensure_logged_in_api_token(page, account: dict, **_kwargs) -> LoginResult:
+    return LoginResult(ok=True, needs_human=False, message="api_token: no browser login required")
+
+
+def _default_sso_hosts(account: dict) -> list[str]:
+    ex = _account_extra(account)
+    hosts = ex.get("sso_redirect_hosts")
+    if isinstance(hosts, list) and hosts:
+        return [str(h) for h in hosts if h]
+    return ["accounts.google.com", "login.microsoftonline.com", "okta.com"]
+
+
+def _url_has_host_fragment(url: str, hosts: list[str]) -> bool:
+    u = (url or "").lower()
+    return any(h.lower() in u for h in hosts)
+
+
+def ensure_logged_in_sso_redirect(
+    page,
+    account: dict,
+    *,
+    wait_human_sec: int = 300,
+    **_kwargs,
+) -> LoginResult:
+    if _is_already_logged_in(page, account):
+        return LoginResult(ok=True, needs_human=False, message="already logged in (profile reused)")
+    url = account.get("url") or ""
+    if not url:
+        return LoginResult(ok=False, needs_human=True, message="account url missing for SSO flow")
+
+    target_netloc = (urlparse(url).netloc or "").lower()
+    page.goto(url, wait_until="domcontentloaded", timeout=60000)
+
+    hosts = _default_sso_hosts(account)
+    seen_sso = _poll_until(lambda: _url_has_host_fragment(page.url or "", hosts), min(60, wait_human_sec), 2.0)
+    if not seen_sso:
+        return LoginResult(ok=False, needs_human=True, message="sso redirect not detected")
+
+    _stderr_prompt("请在弹出的 SSO 页面完成登录")
+
+    def _back_home() -> bool:
+        try:
+            cur = (urlparse(page.url).netloc or "").lower()
+            if target_netloc and cur == target_netloc:
+                return True
+        except Exception:
+            pass
+        return _is_already_logged_in(page, account)
+
+    ok = _poll_until(_back_home, wait_human_sec, interval_sec=2.0)
+    return LoginResult(
+        ok=ok,
+        needs_human=True,
+        message="sso completed" if ok else "sso timeout",
+    )
+
+
+def ensure_logged_in(
+    page,
+    account: dict,
+    *,
+    credentials: dict | None = None,
+    **kwargs,
+) -> LoginResult:
+    """根据 account['auth_strategy'] 分派到对应 helper。"""
+    strat = (account.get("auth_strategy") or "").strip()
+
+    if not strat:
+        raise AuthStrategyError(
+            "ERR_AUTH_STRATEGY_NOT_SUPPORTED",
+            "account has no auth_strategy set",
+        )
+
+    def _need_password() -> str:
+        if not credentials or not credentials.get("plaintext"):
+            raise AuthStrategyError(
+                "ERR_CREDENTIAL_MISSING",
+                f"auth_strategy={strat!r} requires credentials.plaintext",
+            )
+        return str(credentials["plaintext"])
+
+    if strat == "password_auto":
+        return ensure_logged_in_password_auto(page, account, _need_password(), **kwargs)
+    if strat == "password_with_captcha":
+        return ensure_logged_in_password_with_captcha(page, account, _need_password(), **kwargs)
+    if strat == "password_plus_2fa":
+        return ensure_logged_in_password_plus_2fa(page, account, _need_password(), **kwargs)
+    if strat == "qr_code_manual":
+        return ensure_logged_in_qr_code_manual(page, account, **kwargs)
+    if strat == "per_session_manual":
+        return ensure_logged_in_per_session_manual(page, account, **kwargs)
+    if strat == "device_bound":
+        password = None
+        if credentials and credentials.get("plaintext"):
+            password = str(credentials["plaintext"])
+        return ensure_logged_in_device_bound(page, account, password, **kwargs)
+    if strat == "client_certificate":
+        return ensure_logged_in_client_certificate(page, account)
+    if strat == "api_token":
+        return ensure_logged_in_api_token(page, account)
+    if strat == "sso_redirect":
+        return ensure_logged_in_sso_redirect(page, account, **kwargs)
+
+    raise AuthStrategyError(
+        "ERR_AUTH_STRATEGY_NOT_SUPPORTED",
+        f"unknown auth_strategy: {strat!r}",
+    )