tools/export_storage_state.py

"""Export an account-owner-approved Arena Playwright storage_state.

Run this only on your own computer. The generated files contain active session
credentials and must be stored as secrets, never committed to Git.

For Google OAuth, launch a regular installed Chrome with a dedicated profile and
local CDP port, complete login manually, then run this tool with --cdp-url. This
avoids asking Google to authenticate inside Playwright's bundled Chromium.
"""

from __future__ import annotations

import argparse
import base64
import gzip
import json
from pathlib import Path
from urllib.parse import urlparse

from playwright.sync_api import sync_playwright


def find_arena_page(context):
    for page in context.pages:
        if "arena.ai" in page.url:
            return page
    return None


def is_arena_host(host: str) -> bool:
    normalized = host.strip().lower().lstrip(".")
    return normalized == "arena.ai" or normalized.endswith(".arena.ai") or normalized == "lmarena.ai" or normalized.endswith(".lmarena.ai")


def write_state(context, state_path: Path, base64_path: Path) -> None:
    raw_state = context.storage_state()
    # CDP may expose cookies for every site in the dedicated profile. Export only
    # Arena data so Google account cookies are never placed in the Space Secret.
    state = {
        "cookies": [
            cookie
            for cookie in raw_state.get("cookies", [])
            if is_arena_host(str(cookie.get("domain") or ""))
        ],
        "origins": [
            origin
            for origin in raw_state.get("origins", [])
            if is_arena_host(urlparse(str(origin.get("origin") or "")).hostname or "")
        ],
    }
    if not state["cookies"] and not state["origins"]:
        raise RuntimeError(
            "No Arena cookies or Local Storage were found. Complete Arena login and send a "
            "message in the connected browser before exporting."
        )
    state_text = json.dumps(state, ensure_ascii=False, separators=(",", ":"))
    state_path.write_text(state_text, encoding="utf-8")
    base64_path.write_text(
        base64.b64encode(gzip.compress(state_text.encode("utf-8"), compresslevel=9)).decode(
            "ascii"
        ),
        encoding="ascii",
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--state-file", default="arena_storage_state.json")
    parser.add_argument("--base64-file", default="arena_storage_state.b64")
    parser.add_argument(
        "--cdp-url",
        default="",
        help=(
            "Connect to a regular Chrome/Edge launched with a local remote-debugging port, "
            "for example http://127.0.0.1:9222"
        ),
    )
    args = parser.parse_args()

    state_path = Path(args.state_file)
    base64_path = Path(args.base64_file)

    with sync_playwright() as playwright:
        if args.cdp_url:
            browser = playwright.chromium.connect_over_cdp(args.cdp_url)
            if not browser.contexts:
                raise RuntimeError("The CDP browser has no available browser context")
            context = browser.contexts[0]
            page = find_arena_page(context)
            if page is None:
                page = context.new_page()
                page.goto("https://arena.ai/text/direct", wait_until="domcontentloaded")
            page.bring_to_front()
            print("\n已連接到一般 Chrome／Edge。")
        else:
            # Non-Google/manual flow. Google may reject authentication inside this
            # Playwright-managed Chromium; use --cdp-url for Google OAuth.
            browser = playwright.chromium.launch(
                headless=False,
                args=["--start-maximized"],
            )
            context = browser.new_context(no_viewport=True)
            page = context.new_page()
            page.goto("https://arena.ai/text/direct", wait_until="domcontentloaded")

        print("\n請在剛開啟或已連接的官方 Arena 視窗中：")
        print("1. 使用自己的帳號登入。")
        print("2. 自行閱讀並決定是否接受服務條款。")
        print("3. 完成必要的人工驗證。")
        print("4. 實際傳送一則訊息並確認模型能回答。")
        input("\n全部完成後回到這個終端機，按 Enter 匯出 storage_state...")

        write_state(context, state_path, base64_path)
        if not args.cdp_url:
            browser.close()

    print(f"\n已建立：{state_path}")
    print(f"已建立：{base64_path}")
    print("將 .b64 檔案的完整內容設為 Hugging Face Secret ARENA_STORAGE_STATE_B64。")
    print("警告：輸出檔與 CDP 專用瀏覽器 profile 都包含登入 Session。")
    print("設定 Secret 後請安全刪除，不可提交 Git，也不要公開 remote-debugging port。")


if __name__ == "__main__":
    main()