452 lines
14 KiB
Python
452 lines
14 KiB
Python
# -*- 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 提供正向证据（URL 正则或可见 selector）时为 True。"""
|
||
try:
|
||
cur = page.url or ""
|
||
except Exception:
|
||
return False
|
||
cur_stripped = cur.strip()
|
||
if not cur_stripped:
|
||
return False
|
||
ul = cur_stripped.lower()
|
||
if ul == "about:blank" or ul.startswith("about:blank"):
|
||
return False
|
||
|
||
extra = _account_extra(account)
|
||
|
||
pattern = extra.get("logged_in_url_pattern")
|
||
if isinstance(pattern, str) and pattern.strip():
|
||
try:
|
||
if re.search(pattern, cur_stripped):
|
||
return True
|
||
except re.error:
|
||
pass
|
||
|
||
sel = extra.get("logged_in_selector")
|
||
if isinstance(sel, str) and sel.strip():
|
||
try:
|
||
if page.locator(sel.strip()).first.is_visible(timeout=500):
|
||
return True
|
||
except Exception:
|
||
pass
|
||
|
||
return False
|
||
|
||
|
||
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}",
|
||
)
|