diff --git a/scripts/service/rpa_helpers/auth_flows.py b/scripts/service/rpa_helpers/auth_flows.py new file mode 100644 index 0000000..c2bc3c3 --- /dev/null +++ b/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}", + )