AutoRpa/rpa_selection/auth_login.py

294 lines
11 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

from __future__ import annotations
import json
import os
import time
from pathlib import Path
from typing import Any
from http_util import parse_cookie_value
from mall_config import (
YAOEX_HOME_URL,
YAOEX_LOGIN_URL,
YBM_HOME_URL,
YBM_LOGIN_URL,
)
from search_yaoex import YaoexAuth
from util import normalize_text
PROJECT_ROOT = Path(__file__).resolve().parent
AUTH_DIR = PROJECT_ROOT / "data" / ".auth"
DEFAULT_YAOEX_AUTH_PATH = AUTH_DIR / "yaoex.json"
DEFAULT_YBM_COOKIE_PATH = AUTH_DIR / "ybm.json"
def load_dotenv(path: Path) -> None:
if not path.exists():
return
for line in path.read_text(encoding="utf-8").splitlines():
stripped = line.strip()
if not stripped or stripped.startswith("#") or "=" not in stripped:
continue
key, _, value = stripped.partition("=")
key = key.strip()
value = value.strip().strip('"').strip("'")
os.environ.setdefault(key, value)
def load_secrets(secrets_path: Path | None = None) -> dict[str, dict[str, str]]:
load_dotenv(PROJECT_ROOT / ".env")
result: dict[str, dict[str, str]] = {"yaoex": {}, "ybm": {}}
path = secrets_path or (PROJECT_ROOT / "secrets.json")
if path.exists():
data = json.loads(path.read_text(encoding="utf-8"))
if isinstance(data, dict):
for mall in ("yaoex", "ybm"):
block = data.get(mall) or {}
if isinstance(block, dict):
result[mall] = {
"username": normalize_text(block.get("username")),
"password": normalize_text(block.get("password")),
}
env_map = {
"yaoex": ("YAOEX_USERNAME", "YAOEX_PASSWORD"),
"ybm": ("YBM_USERNAME", "YBM_PASSWORD"),
}
for mall, (user_key, pass_key) in env_map.items():
result[mall]["username"] = os.environ.get(user_key) or result[mall].get("username", "")
result[mall]["password"] = os.environ.get(pass_key) or result[mall].get("password", "")
return result
def load_yaoex_auth_file(path: Path) -> YaoexAuth:
data = json.loads(path.read_text(encoding="utf-8"))
return YaoexAuth(
token=normalize_text(data.get("token")),
user_id=normalize_text(data.get("user_id")),
buyer_code=normalize_text(data.get("buyer_code")),
)
def load_cookie_file(path: Path) -> dict[str, str]:
data = json.loads(path.read_text(encoding="utf-8"))
if not isinstance(data, dict):
raise ValueError(f"cookie 文件不是对象:{path}")
return {str(key): parse_cookie_value(str(value)) for key, value in data.items()}
def save_yaoex_auth(auth: YaoexAuth, path: Path = DEFAULT_YAOEX_AUTH_PATH) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(
json.dumps(
{"token": auth.token, "user_id": auth.user_id, "buyer_code": auth.buyer_code},
ensure_ascii=False,
indent=2,
),
encoding="utf-8",
)
def save_cookie_file(cookies: dict[str, str], path: Path = DEFAULT_YBM_COOKIE_PATH) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(cookies, ensure_ascii=False, indent=2), encoding="utf-8")
def _require_account(mall_label: str, username: str, password: str) -> None:
if not username or not password:
raise SystemExit(
f"缺少{mall_label}账号。请在 .env / 环境变量或 secrets.json 中配置用户名和密码。"
)
class BrowserRuntime:
"""同一轮脚本共用一套 Playwright。keep=True 时搜完仍留着窗口。"""
def __init__(self, keep: bool = False, hold_seconds: int = 180) -> None:
self.keep = keep
self.hold_seconds = hold_seconds
self._playwright: Any = None
self._browsers: list[Any] = []
def start(self) -> None:
if self._playwright is not None:
return
from playwright.sync_api import sync_playwright
self._playwright = sync_playwright().start()
def new_page(self, *, viewport: dict[str, int] | None = None) -> tuple[Any, Any]:
self.start()
browser = self._playwright.chromium.launch(
headless=False,
channel="chrome",
args=["--new-window", "--start-maximized"],
)
self._browsers.append(browser)
context = browser.new_context(viewport=viewport or {"width": 1440, "height": 900})
page = context.new_page()
page.bring_to_front()
return context, page
def close(self) -> None:
if self.keep and self._browsers:
if self.hold_seconds <= 0:
print("浏览器保持打开,不会自动关。要关时在终端停掉进程。")
while True:
time.sleep(3600)
else:
print(f"浏览器保持打开 {self.hold_seconds} 秒,请看屏幕上的 Chrome。")
time.sleep(self.hold_seconds)
for browser in self._browsers:
try:
browser.close()
except Exception:
pass
self._browsers.clear()
if self._playwright is not None:
try:
self._playwright.stop()
except Exception:
pass
self._playwright = None
def _ensure_yaoex_logged_in(page: Any, username: str, password: str) -> None:
page.goto(YAOEX_LOGIN_URL, wait_until="domcontentloaded")
page.wait_for_timeout(2000)
if "/login" not in page.url:
return
account_input = page.locator("#username").first
password_input = page.locator("#password").first
captcha_button = page.locator(
"#captcha-box .geetest_btn_click, #captcha-box .geetest_btn, #captcha-box .geetest_radar_btn"
).first
login_button = page.locator("#login-btn").first
account_input.wait_for(state="visible", timeout=30000)
account_input.fill(username)
password_input.fill(password)
print("1号药城:已填账号。请看 Chrome——若出现滑块拼图,拖过去。窗口会等 3 分钟。")
if captcha_button.count() > 0:
captcha_button.click(force=True)
try:
page.wait_for_function(
"() => !window.location.href.includes('/login')",
timeout=20000,
)
except Exception:
try:
login_button.click(force=True)
except Exception:
pass
try:
page.wait_for_function(
"() => !window.location.href.includes('/login')",
timeout=180000,
)
except Exception:
pass
if "/login" in page.url:
login_error = ""
error_locator = page.locator(".error-tips span, .layui-layer-content, .geetest_tip")
if error_locator.count() > 0:
login_error = normalize_text(error_locator.first.inner_text())
shot = PROJECT_ROOT / "output" / "结果" / "yaoex_login_fail.png"
shot.parent.mkdir(parents=True, exist_ok=True)
page.screenshot(path=str(shot), full_page=True)
detail = f",页面提示:{login_error}" if login_error else ""
raise RuntimeError(
f"1号药城登录后仍停留在登录页,登录可能失败{detail}。截图:{shot}"
)
def _ensure_ybm_logged_in(page: Any, username: str, password: str) -> None:
page.goto(YBM_LOGIN_URL, wait_until="domcontentloaded")
page.wait_for_timeout(2000)
if "/new/login" not in page.url:
return
account_input = page.locator("form.login-form input[type='text']").first
password_input = page.locator("form.login-form input[type='password']").first
agree_input = page.locator(".login-section-agree input[type='checkbox']").first
agree_visible_box = page.locator(".login-section-agree .el-checkbox__inner").first
login_button = page.locator("button.login-button").first
account_input.wait_for(state="visible", timeout=30000)
account_input.fill(username)
password_input.fill(password)
if not agree_input.is_checked():
agree_visible_box.click()
login_button.click()
page.wait_for_timeout(5000)
if "/new/login" in page.url:
login_error = ""
error_locator = page.locator(".el-message__content")
if error_locator.count() > 0:
login_error = normalize_text(error_locator.first.inner_text())
detail = f",页面提示:{login_error}" if login_error else ""
raise RuntimeError(f"药帮忙登录后仍停留在登录页,登录可能失败{detail}")
def _playwright_login_yaoex(username: str, password: str, runtime: BrowserRuntime) -> YaoexAuth:
context, page = runtime.new_page()
_ensure_yaoex_logged_in(page, username, password)
page.goto(YAOEX_HOME_URL, wait_until="domcontentloaded")
page.wait_for_timeout(3000)
print("1号药城窗口应已打开首页。")
cookies = context.cookies()
cookie_map = {item["name"]: item["value"] for item in cookies}
token = parse_cookie_value(cookie_map.get("ycgltoken")) or parse_cookie_value(
cookie_map.get("yctoken")
)
user_id = parse_cookie_value(cookie_map.get("ycuserId"))
buyer_code = parse_cookie_value(cookie_map.get("ycenterpriseId"))
if not token or not user_id or not buyer_code:
raise RuntimeError("未能从1号药城 cookie 提取完整认证字段(token/user_id/buyer_code)。")
return YaoexAuth(token=token, user_id=user_id, buyer_code=buyer_code)
def _playwright_login_ybm(username: str, password: str, runtime: BrowserRuntime) -> dict[str, str]:
context, page = runtime.new_page()
_ensure_ybm_logged_in(page, username, password)
page.goto(YBM_HOME_URL, wait_until="domcontentloaded")
page.wait_for_timeout(3000)
print("药帮忙窗口应已打开首页。")
cookies = context.cookies()
return {item["name"]: parse_cookie_value(item["value"]) for item in cookies}
def login_yaoex(
secrets_path: Path | None = None,
runtime: BrowserRuntime | None = None,
) -> YaoexAuth:
secrets = load_secrets(secrets_path)
username = secrets["yaoex"].get("username", "")
password = secrets["yaoex"].get("password", "")
_require_account("1号药城", username, password)
own = runtime is None
session = runtime or BrowserRuntime(keep=False)
try:
return _playwright_login_yaoex(username, password, session)
finally:
if own:
session.close()
def login_ybm(
secrets_path: Path | None = None,
runtime: BrowserRuntime | None = None,
) -> dict[str, str]:
secrets = load_secrets(secrets_path)
username = secrets["ybm"].get("username", "")
password = secrets["ybm"].get("password", "")
_require_account("药帮忙", username, password)
own = runtime is None
session = runtime or BrowserRuntime(keep=False)
try:
return _playwright_login_ybm(username, password, session)
finally:
if own:
session.close()