feat(rpa_helpers): add auth_flows.py with 9 ensure_logged_in_* and dispatcher
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
449
scripts/service/rpa_helpers/auth_flows.py
Normal file
449
scripts/service/rpa_helpers/auth_flows.py
Normal file
@@ -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}",
|
||||
)
|
||||
Reference in New Issue
Block a user