368 lines
14 KiB
Python
368 lines
14 KiB
Python
# 商城登录：读环境变量账号，写出 token/cookie 缓存。
|
||
#
|
||
# 账号只认 YAOEX_ / YBM_ 环境变量，不读 secrets.json。
|
||
# 缓存文件权限收成 0o600。
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import sys
|
||
import time
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from http_util import cookies_for_http, 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
|
||
REPO_ROOT = PROJECT_ROOT.parent
|
||
if str(REPO_ROOT) not in sys.path:
|
||
sys.path.insert(0, str(REPO_ROOT))
|
||
|
||
from rpa_oracle.env import load_env
|
||
|
||
AUTH_DIR = PROJECT_ROOT / "data" / ".auth"
|
||
DEFAULT_YAOEX_AUTH_PATH = AUTH_DIR / "yaoex.json"
|
||
DEFAULT_YBM_COOKIE_PATH = AUTH_DIR / "ybm.json"
|
||
|
||
|
||
# 把单个 .env 补进环境。已有非空变量不覆盖。加购班过渡仍走这一层。
|
||
def load_dotenv(path: Path) -> None:
|
||
load_env(path)
|
||
|
||
|
||
# 读商城账号。只认 YAOEX_ / YBM_ 环境变量，不读 secrets.json。
|
||
#
|
||
# 先补选品 .env 里尚未占用的键。
|
||
def load_secrets() -> dict[str, dict[str, str]]:
|
||
load_env(PROJECT_ROOT / ".env")
|
||
env_map = {
|
||
"yaoex": ("YAOEX_USERNAME", "YAOEX_PASSWORD"),
|
||
"ybm": ("YBM_USERNAME", "YBM_PASSWORD"),
|
||
}
|
||
return {
|
||
mall: {
|
||
"username": os.environ.get(user_key) or "",
|
||
"password": os.environ.get(pass_key) or "",
|
||
}
|
||
for mall, (user_key, pass_key) in env_map.items()
|
||
}
|
||
|
||
|
||
# 从缓存 JSON 读出 1 号药城 token / user_id / buyer_code。
|
||
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")),
|
||
)
|
||
|
||
|
||
# 从缓存 JSON 读出药帮忙 cookie。文件必须是对象。只留下能进 HTTP 头的项。
|
||
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}")
|
||
loaded = {str(key): parse_cookie_value(str(value)) for key, value in data.items()}
|
||
return cookies_for_http(loaded)
|
||
|
||
|
||
# 把缓存文件权限收成 0o600。Windows 上 chmod 可能无效，忽略 OSError。
|
||
def _restrict_file_mode(path: Path) -> None:
|
||
try:
|
||
os.chmod(path, 0o600)
|
||
except OSError:
|
||
pass
|
||
|
||
|
||
# 写出 1 号药城登录缓存，权限 0o600。
|
||
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",
|
||
)
|
||
_restrict_file_mode(path)
|
||
|
||
|
||
# 写出药帮忙 cookie 缓存，权限 0o600。中文展示字段不写进文件。
|
||
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_for_http(cookies), ensure_ascii=False, indent=2),
|
||
encoding="utf-8",
|
||
)
|
||
_restrict_file_mode(path)
|
||
|
||
|
||
# 账号或密码为空时退出，提示去环境变量配置。
|
||
def _require_account(mall_label: str, username: str, password: str) -> None:
|
||
if not username or not password:
|
||
raise SystemExit(
|
||
f"缺少{mall_label}账号。请在 .env 或环境变量中配置用户名和密码。"
|
||
)
|
||
|
||
|
||
# 路径是否落在仓库根之内。
|
||
def _is_inside_repo(path: Path, repo_root: Path) -> bool:
|
||
try:
|
||
path.resolve().relative_to(repo_root.resolve())
|
||
except ValueError:
|
||
return False
|
||
return True
|
||
|
||
|
||
# 选品 Chrome 用户目录固定在仓库外，避免每次新开空浏览器触发极验滑块。
|
||
def resolve_selection_chrome_profile(
|
||
requested: str | None = None,
|
||
*,
|
||
repo_root: Path | None = None,
|
||
) -> Path:
|
||
root = (repo_root or REPO_ROOT).resolve()
|
||
outside = (root.parent / "rpa-chrome-profile" / "selection").resolve()
|
||
raw = os.getenv("RPA_SELECTION_CHROME_PROFILE") if requested is None else requested
|
||
if not raw:
|
||
base = os.getenv("RPA_CHROME_PROFILE")
|
||
raw = str(Path(base).expanduser() / "selection") if base else ""
|
||
if not raw:
|
||
return outside
|
||
candidate = Path(raw).expanduser().resolve()
|
||
if _is_inside_repo(candidate, root):
|
||
print(f"选品 Chrome 用户目录指向仓库内 {candidate}，已改用 {outside}")
|
||
return outside
|
||
return candidate
|
||
|
||
|
||
# 同一轮脚本共用一套已登录的 Chrome 用户目录。keep=True 时搜完仍留着窗口。
|
||
class BrowserRuntime:
|
||
# 记下是否保持窗口，以及保持多少秒。
|
||
def __init__(self, keep: bool = False, hold_seconds: int = 180) -> None:
|
||
self.keep = keep
|
||
self.hold_seconds = hold_seconds
|
||
self._playwright: Any = None
|
||
self._context: Any = None
|
||
self._virgin_page: Any = None
|
||
|
||
# 打开仓库外的持久 Chrome 用户目录。已启动则直接返回。
|
||
def start(self, *, viewport: dict[str, int] | None = None) -> None:
|
||
if self._playwright is not None:
|
||
return
|
||
from playwright.sync_api import sync_playwright
|
||
|
||
self._playwright = sync_playwright().start()
|
||
profile = resolve_selection_chrome_profile()
|
||
profile.mkdir(parents=True, exist_ok=True)
|
||
print(f"选品 Chrome 用户目录：{profile}")
|
||
try:
|
||
self._context = self._playwright.chromium.launch_persistent_context(
|
||
str(profile),
|
||
headless=False,
|
||
channel="chrome",
|
||
viewport=viewport or {"width": 1440, "height": 900},
|
||
args=[
|
||
"--start-maximized",
|
||
"--disable-blink-features=AutomationControlled",
|
||
"--no-first-run",
|
||
"--no-default-browser-check",
|
||
],
|
||
ignore_default_args=["--enable-automation"],
|
||
)
|
||
except Exception as exc:
|
||
raise RuntimeError(
|
||
f"无法打开选品 Chrome 用户目录 {profile}（是否已有实例占用）：{exc}"
|
||
) from exc
|
||
pages = list(self._context.pages)
|
||
self._virgin_page = pages[0] if pages else None
|
||
|
||
# 取出一个页面。优先复用启动时自带的空白页。
|
||
def new_page(self, *, viewport: dict[str, int] | None = None) -> tuple[Any, Any]:
|
||
self.start(viewport=viewport)
|
||
assert self._context is not None
|
||
if self._virgin_page is not None:
|
||
page = self._virgin_page
|
||
self._virgin_page = None
|
||
else:
|
||
page = self._context.new_page()
|
||
page.bring_to_front()
|
||
return self._context, page
|
||
|
||
# keep=True 时先留窗口，再关 Playwright。
|
||
def close(self) -> None:
|
||
if self.keep and self._context is not None:
|
||
if self.hold_seconds <= 0:
|
||
print("浏览器保持打开，不会自动关。要关时在终端停掉进程。")
|
||
while True:
|
||
time.sleep(3600)
|
||
else:
|
||
print(f"浏览器保持打开 {self.hold_seconds} 秒，请看屏幕上的 Chrome。")
|
||
time.sleep(self.hold_seconds)
|
||
if self._context is not None:
|
||
try:
|
||
self._context.close()
|
||
except Exception:
|
||
pass
|
||
self._context = None
|
||
self._virgin_page = None
|
||
if self._playwright is not None:
|
||
try:
|
||
self._playwright.stop()
|
||
except Exception:
|
||
pass
|
||
self._playwright = None
|
||
|
||
|
||
# 若停在 1 号药城登录页则填账号；仍停在登录页则失败。
|
||
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}")
|
||
|
||
|
||
# 打开 1 号药城，登录后从 cookie 抽出 token / user_id / buyer_code。
|
||
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)
|
||
|
||
|
||
# 打开药帮忙，登录后带回能进 HTTP 头的 cookie。
|
||
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()
|
||
loaded = {item["name"]: parse_cookie_value(item["value"]) for item in cookies}
|
||
return cookies_for_http(loaded)
|
||
|
||
|
||
# 登录 1 号药城。缺账号则退出；传入 runtime 时不关浏览器。
|
||
def login_yaoex(runtime: BrowserRuntime | None = None) -> YaoexAuth:
|
||
secrets = load_secrets()
|
||
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()
|
||
|
||
|
||
# 登录药帮忙。缺账号则退出；传入 runtime 时不关浏览器。
|
||
def login_ybm(runtime: BrowserRuntime | None = None) -> dict[str, str]:
|
||
secrets = load_secrets()
|
||
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()
|