From 5fada75d236ed79e1017679844089d4f86437f5b Mon Sep 17 00:00:00 2001 From: Yuxuan Yan Date: Sat, 4 Jul 2026 18:37:34 +0800 Subject: [PATCH] Add JoinQuant env browser login helper --- plugins/joinquant/browser.py | 93 +++++++++++++++++++++++++++++++++- plugins/joinquant/cli.py | 50 ++++++++++++++++++ tests/test_joinquant_plugin.py | 15 ++++++ 3 files changed, 157 insertions(+), 1 deletion(-) diff --git a/plugins/joinquant/browser.py b/plugins/joinquant/browser.py index 034317b..4ad98a3 100644 --- a/plugins/joinquant/browser.py +++ b/plugins/joinquant/browser.py @@ -188,6 +188,23 @@ def load_json(path: str | Path) -> dict[str, Any]: return data +def load_env_file(path: str | Path) -> dict[str, str]: + """Load simple KEY=VALUE env files without requiring shell-safe quoting.""" + values: dict[str, str] = {} + for raw in Path(path).expanduser().read_text(encoding="utf-8").splitlines(): + line = raw.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, value = line.split("=", 1) + clean = value.strip() + if len(clean) >= 2 and clean[0] == clean[-1] and clean[0] in {"'", '"'}: + clean = clean[1:-1] + else: + clean = clean.strip("'\"") + values[key.strip()] = clean + return values + + def _get_dotted(data: dict[str, Any], dotted: str) -> Any: current: Any = data for part in dotted.split("."): @@ -271,6 +288,80 @@ def save_login_state( return state_path +def save_login_state_from_env( + *, + env_path: str | Path, + storage_state: str | Path, + login_url: str = DEFAULT_LOGIN_URL, + headless: bool = True, + out_dir: str | Path | None = None, + timeout_ms: int = 120_000, +) -> dict[str, Any]: + """Try to log in with env-file credentials and save browser state. + + This handles the normal password-login form. If JoinQuant presents CAPTCHA, + slide verification, SMS verification, or 2FA, the report will mark the run + as not logged in and save a screenshot for manual diagnosis. + """ + env = load_env_file(env_path) + username = env.get("JOINQUANT_USERNAME") + password = env.get("JOINQUANT_PASSWORD") + if not username or not password: + raise ValueError("JOINQUANT_USERNAME and JOINQUANT_PASSWORD are required") + + sync_playwright, _ = _require_playwright() + state_path = Path(storage_state).expanduser() + state_path.parent.mkdir(parents=True, exist_ok=True) + artifact_dir = Path(out_dir or state_path.parent / "joinquant_login_artifacts") + artifact_dir.mkdir(parents=True, exist_ok=True) + screenshot_path = artifact_dir / "login_after_submit.png" + html_path = artifact_dir / "login_after_submit.html" + + with sync_playwright() as pw: + browser = pw.chromium.launch(headless=headless) + context = browser.new_context() + page = context.new_page() + page.set_default_timeout(timeout_ms) + page.goto(login_url, wait_until="networkidle", timeout=timeout_ms) + page.locator('input[name="username"], input.pwd-phone').first.fill(username) + page.locator('input[name="pwd"], input.jq-login__password').first.fill(password) + checkbox = page.locator("#agreementBox, input.agreement-box").first + if checkbox.count(): + checkbox.check(force=True) + page.locator("button.btnPwdSubmit, button.login-submit").first.click() + page.wait_for_timeout(5000) + + html = page.content() + html_path.write_text(html, encoding="utf-8") + page.screenshot(path=str(screenshot_path), full_page=True) + login_inputs = page.locator('input[name="username"], input[name="pwd"]').count() + current_url = page.url + logged_in = login_inputs == 0 and "/user/login" not in current_url + if logged_in: + context.storage_state(path=str(state_path)) + state_path.chmod(0o600) + browser.close() + + report = { + "created_at": datetime.now(timezone.utc).isoformat(), + "logged_in": bool(logged_in), + "storage_state": str(state_path) if logged_in else "", + "artifact_dir": str(artifact_dir), + "screenshot": str(screenshot_path), + "html": str(html_path), + "current_url": current_url, + "notes": ( + "Logged in and saved browser state." + if logged_in + else "Login did not complete. CAPTCHA/SMS/2FA or invalid credentials may be required." + ), + } + report_path = artifact_dir / "login_report.json" + report["report_path"] = str(report_path) + report_path.write_text(json.dumps(report, indent=2, ensure_ascii=False) + "\n", encoding="utf-8") + return report + + def browser_snapshot( *, url: str, @@ -299,7 +390,7 @@ def browser_snapshot( def _locator(page: Any, selector: str): - return page.locator(selector).first() + return page.locator(selector).first def _action_fail(action: dict[str, Any], exc: Exception) -> None: diff --git a/plugins/joinquant/cli.py b/plugins/joinquant/cli.py index 6554de1..6eb96fb 100644 --- a/plugins/joinquant/cli.py +++ b/plugins/joinquant/cli.py @@ -6,9 +6,11 @@ import click from plugins.joinquant.browser import ( browser_snapshot, + load_env_file, run_browser_backtest, run_browser_sim_trade, save_login_state, + save_login_state_from_env, write_browser_config_template, ) from plugins.joinquant.export_targets import export_targets @@ -257,6 +259,54 @@ def browser_login_cmd(storage_state, login_url, headless, wait_seconds): click.echo(f"Saved JoinQuant browser state: {path}") +@joinquant.command("browser-login-env") +@click.option( + "--env-path", + default="~/.config/chinese-equity-quant/joinquant.env", + show_default=True, + help="Env file with JOINQUANT_USERNAME and JOINQUANT_PASSWORD", +) +@click.option( + "--storage-state", + default="~/.config/chinese-equity-quant/joinquant_storage_state.json", + show_default=True, + help="Path where the authenticated browser state is stored", +) +@click.option( + "--login-url", + default="https://www.joinquant.com/user/login/index", + show_default=True, +) +@click.option("--headed", is_flag=True, help="Run with a visible browser") +@click.option("--out-dir", default=None, help="Login artifact directory") +@click.option("--timeout-ms", default=120_000, show_default=True, type=int) +def browser_login_env_cmd(env_path, storage_state, login_url, headed, out_dir, timeout_ms): + """Try credential-based JoinQuant login from an env file.""" + env = load_env_file(env_path) + missing = [ + key for key in ["JOINQUANT_USERNAME", "JOINQUANT_PASSWORD"] + if not env.get(key) + ] + if missing: + raise click.ClickException(f"Missing required env keys: {', '.join(missing)}") + report = save_login_state_from_env( + env_path=env_path, + storage_state=storage_state, + login_url=login_url, + headless=not headed, + out_dir=out_dir, + timeout_ms=timeout_ms, + ) + click.echo(f"Saved login report: {report['report_path']}") + if report["logged_in"]: + click.echo(f"Saved JoinQuant browser state: {report['storage_state']}") + else: + raise click.ClickException( + "JoinQuant login did not complete. Check the saved screenshot/HTML; " + "CAPTCHA/SMS/2FA may require `joinquant browser-login` once." + ) + + @joinquant.command("write-browser-config") @click.option("--out-path", required=True, help="Path for JSON browser automation config") @click.option("--strategy-url", default="", help="JoinQuant strategy edit/backtest/模拟盘 URL") diff --git a/tests/test_joinquant_plugin.py b/tests/test_joinquant_plugin.py index 55b8da1..8309be8 100644 --- a/tests/test_joinquant_plugin.py +++ b/tests/test_joinquant_plugin.py @@ -14,6 +14,7 @@ from cli import cli from pipeline.common.schema import FILL_COLUMNS, PNL_COLUMNS, POSITION_COLUMNS from plugins.joinquant.browser import ( default_browser_config, + load_env_file, resolve_template, write_browser_config_template, ) @@ -559,6 +560,20 @@ def test_browser_config_template_and_placeholder_resolution(tmp_path): assert resolve_template("save:{expected_joinquant_csvs.fills}", context) == "save:/tmp/jq_fills.csv" +def test_load_env_file_handles_quotes_without_shell_sourcing(tmp_path): + env_path = tmp_path / "joinquant.env" + env_path.write_text( + "JOINQUANT_USERNAME=alice\n" + "JOINQUANT_PASSWORD=\"secret\"\n" + "JOINQUANT_STRATEGY_URL='https://example.test/path?x=1&y=2'\n" + ) + + env = load_env_file(env_path) + assert env["JOINQUANT_USERNAME"] == "alice" + assert env["JOINQUANT_PASSWORD"] == "secret" + assert env["JOINQUANT_STRATEGY_URL"] == "https://example.test/path?x=1&y=2" + + def test_sim_trade_browser_config_template(tmp_path): config_path = write_browser_config_template( tmp_path / "sim_config.json",