chaos
cd4a2351ae
Enable Chromium on machines without DISPLAY: Python auto headless with Xvfb hints, and start.sh installs/starts Xvfb (xvfb-run or persistent :99) so registration can run headed under a virtual display.
527 lines
19 KiB
Python
527 lines
19 KiB
Python
"""Register-machine hook: mint CPA xai auth after successful registration.
|
|
|
|
OIDC package lives at ./cpa_xai (bundled with this project).
|
|
Optional override: config `api_reverse_tools` / env `API_REVERSE_TOOLS`
|
|
points at a directory that *contains* the `cpa_xai` package.
|
|
|
|
After mint, optionally:
|
|
- copy into local CPA auth-dir (`cpa_copy_to_hotload`)
|
|
- POST into remote CPA Management API (`cpa_auto_import_remote`)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import shutil
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Callable
|
|
from urllib.parse import quote, urlparse, urlunparse
|
|
|
|
_REG_DIR = Path(__file__).resolve().parent
|
|
_DEFAULT_OUT = _REG_DIR / "cpa_auths"
|
|
_DEFAULT_CPA = Path("") # empty = do not assume a machine-local CPA path
|
|
_DOTENV_LOADED = False
|
|
|
|
|
|
def _load_dotenv_once() -> None:
|
|
"""Load project .env into os.environ once (does not override existing env)."""
|
|
global _DOTENV_LOADED
|
|
if _DOTENV_LOADED:
|
|
return
|
|
_DOTENV_LOADED = True
|
|
env_path = _REG_DIR / ".env"
|
|
if not env_path.is_file():
|
|
return
|
|
try:
|
|
for raw in env_path.read_text(encoding="utf-8").splitlines():
|
|
line = raw.strip()
|
|
if not line or line.startswith("#") or "=" not in line:
|
|
continue
|
|
key, _, val = line.partition("=")
|
|
key = key.strip()
|
|
val = val.strip().strip('"').strip("'")
|
|
if key and key not in os.environ:
|
|
os.environ[key] = val
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
|
|
def _ensure_cpa_xai_on_path(tools_dir: str | Path | None = None) -> Path:
|
|
"""Put the parent of `cpa_xai` on sys.path. Default: this project root."""
|
|
if tools_dir:
|
|
tools = Path(tools_dir).expanduser().resolve()
|
|
else:
|
|
env = (os.environ.get("API_REVERSE_TOOLS") or "").strip()
|
|
tools = Path(env).expanduser().resolve() if env else _REG_DIR
|
|
# If user pointed at .../cpa_xai itself, use its parent
|
|
if tools.name == "cpa_xai" and (tools / "__init__.py").is_file():
|
|
tools = tools.parent
|
|
if str(tools) not in sys.path:
|
|
sys.path.insert(0, str(tools))
|
|
return tools
|
|
|
|
|
|
def export_cookies_from_page(page: Any) -> list[dict]:
|
|
"""Best-effort export of cookies from a DrissionPage tab/browser."""
|
|
if page is None:
|
|
return []
|
|
cookies = None
|
|
for getter in (
|
|
lambda: page.cookies(all_domains=True, all_info=True),
|
|
lambda: page.cookies(all_domains=True),
|
|
lambda: page.cookies(),
|
|
):
|
|
try:
|
|
cookies = getter()
|
|
if cookies:
|
|
break
|
|
except TypeError:
|
|
continue
|
|
except Exception:
|
|
continue
|
|
if not cookies:
|
|
try:
|
|
browser = getattr(page, "browser", None)
|
|
if browser is not None:
|
|
cookies = browser.cookies()
|
|
except Exception:
|
|
cookies = None
|
|
if isinstance(cookies, list):
|
|
return [c for c in cookies if isinstance(c, dict)]
|
|
return []
|
|
|
|
|
|
def _normalize_cpa_remote_base(base: str) -> str:
|
|
"""Normalize user-provided CPA root URL to scheme://host[:port].
|
|
|
|
Accepts:
|
|
http://127.0.0.1:8317
|
|
http://host:8317/
|
|
http://host:8317/v0/management
|
|
http://host:8317/v0/management/auth-files
|
|
"""
|
|
raw = (base or "").strip()
|
|
if not raw:
|
|
return ""
|
|
if "://" not in raw:
|
|
raw = "http://" + raw
|
|
parsed = urlparse(raw)
|
|
if not parsed.scheme or not parsed.netloc:
|
|
return raw.rstrip("/")
|
|
# Keep only scheme + netloc (drop path/query/fragment)
|
|
return urlunparse((parsed.scheme, parsed.netloc, "", "", "", "")).rstrip("/")
|
|
|
|
|
|
def _env_flag(name: str) -> bool | None:
|
|
"""Parse env truthy/falsey flag. None if unset/empty."""
|
|
raw = (os.environ.get(name) or "").strip().lower()
|
|
if not raw:
|
|
return None
|
|
if raw in ("1", "true", "yes", "on", "y"):
|
|
return True
|
|
if raw in ("0", "false", "no", "off", "n"):
|
|
return False
|
|
return None
|
|
|
|
|
|
def resolve_cpa_remote_settings(cfg: dict | None = None) -> dict[str, Any]:
|
|
"""Resolve remote CPA import settings from config + env.
|
|
|
|
Enable flag priority:
|
|
1) env CPA_AUTO_IMPORT_REMOTE / CPA_AUTO_IMPORT (if set)
|
|
2) config cpa_auto_import_remote
|
|
3) default false
|
|
|
|
Base/password: config first, then env fallbacks.
|
|
"""
|
|
_load_dotenv_once()
|
|
cfg = cfg or {}
|
|
env_enabled = _env_flag("CPA_AUTO_IMPORT_REMOTE")
|
|
if env_enabled is None:
|
|
env_enabled = _env_flag("CPA_AUTO_IMPORT")
|
|
if env_enabled is not None:
|
|
enabled = env_enabled
|
|
else:
|
|
enabled = bool(cfg.get("cpa_auto_import_remote", False))
|
|
base = (
|
|
(cfg.get("cpa_remote_base") or "").strip()
|
|
or (os.environ.get("CPA_REMOTE_BASE") or "").strip()
|
|
)
|
|
password = (
|
|
(cfg.get("cpa_remote_password") or "").strip()
|
|
or (os.environ.get("CPA_REMOTE_PASSWORD") or "").strip()
|
|
or (os.environ.get("CPA_MANAGEMENT_PASSWORD") or "").strip()
|
|
or (os.environ.get("MANAGEMENT_PASSWORD") or "").strip()
|
|
)
|
|
retries = int(cfg.get("cpa_remote_import_retries", 3) or 3)
|
|
delay = float(cfg.get("cpa_remote_import_retry_delay", 2) or 2)
|
|
return {
|
|
"enabled": enabled,
|
|
"base": _normalize_cpa_remote_base(base),
|
|
"password": password,
|
|
"retries": max(1, retries),
|
|
"delay": max(0.0, delay),
|
|
}
|
|
|
|
|
|
def import_cpa_auth_to_remote(
|
|
path: str | Path,
|
|
*,
|
|
base: str,
|
|
password: str,
|
|
retries: int = 3,
|
|
delay: float = 2.0,
|
|
log: Callable[[str], None] | None = None,
|
|
timeout: float = 12.0,
|
|
) -> dict[str, Any]:
|
|
"""Upload one local xai-*.json to CPA Management API.
|
|
|
|
POST {base}/v0/management/auth-files?name=<filename>
|
|
Authorization: Bearer <password>
|
|
Body: raw JSON file content
|
|
|
|
Failures are returned in the result dict; caller decides whether to ignore.
|
|
"""
|
|
_log = log or (lambda m: print(m, flush=True))
|
|
src = Path(path)
|
|
root = _normalize_cpa_remote_base(base)
|
|
if not root:
|
|
return {"ok": False, "error": "empty cpa_remote_base"}
|
|
if not password:
|
|
return {"ok": False, "error": "empty cpa_remote_password"}
|
|
if not src.is_file():
|
|
return {"ok": False, "error": f"auth file not found: {src}"}
|
|
|
|
name = src.name
|
|
if not name.endswith(".json"):
|
|
return {"ok": False, "error": f"filename must end with .json: {name}"}
|
|
|
|
try:
|
|
body = src.read_bytes()
|
|
except Exception as e: # noqa: BLE001
|
|
return {"ok": False, "error": f"read file: {e}"}
|
|
|
|
url = f"{root}/v0/management/auth-files?name={quote(name)}"
|
|
headers = {
|
|
"Authorization": f"Bearer {password}",
|
|
"X-Management-Key": password,
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
try:
|
|
from curl_cffi import requests as _req # type: ignore
|
|
except Exception as e: # noqa: BLE001
|
|
return {"ok": False, "error": f"import curl_cffi failed: {e}"}
|
|
|
|
last_err = "unknown"
|
|
attempts = max(1, int(retries or 1))
|
|
for i in range(1, attempts + 1):
|
|
try:
|
|
# Admin API is local/remote management — do not use outbound proxy.
|
|
resp = _req.post(
|
|
url,
|
|
data=body,
|
|
headers=headers,
|
|
timeout=timeout,
|
|
proxies={},
|
|
impersonate="chrome",
|
|
)
|
|
status = int(getattr(resp, "status_code", 0) or 0)
|
|
text = ""
|
|
try:
|
|
text = (resp.text or "")[:500]
|
|
except Exception:
|
|
text = ""
|
|
if status in (200, 207):
|
|
# Prefer JSON status field when present
|
|
ok_status = True
|
|
try:
|
|
j = resp.json()
|
|
st = str((j or {}).get("status") or "").lower()
|
|
if st and st not in ("ok", "partial", "success"):
|
|
ok_status = False
|
|
last_err = f"status={st} body={text}"
|
|
except Exception:
|
|
pass
|
|
if ok_status:
|
|
_log(f"[cpa] remote import ok -> {root} name={name}")
|
|
return {
|
|
"ok": True,
|
|
"base": root,
|
|
"name": name,
|
|
"status_code": status,
|
|
"body": text,
|
|
}
|
|
else:
|
|
last_err = f"HTTP {status}: {text or '(empty)'}"
|
|
except Exception as e: # noqa: BLE001
|
|
last_err = str(e)
|
|
|
|
if i < attempts:
|
|
_log(f"[cpa] remote import retry {i}/{attempts}: {last_err}")
|
|
if delay > 0:
|
|
time.sleep(delay)
|
|
|
|
_log(f"[cpa] remote import failed: {last_err}")
|
|
return {
|
|
"ok": False,
|
|
"base": root,
|
|
"name": name,
|
|
"error": last_err,
|
|
}
|
|
|
|
|
|
def maybe_import_cpa_auth_remote(
|
|
path: str | Path,
|
|
*,
|
|
config: dict | None = None,
|
|
log_callback: Callable[[str], None] | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Config-gated remote import helper (safe no-op when disabled)."""
|
|
cfg = config or {}
|
|
log = log_callback or (lambda m: print(m, flush=True))
|
|
settings = resolve_cpa_remote_settings(cfg)
|
|
if not settings["enabled"]:
|
|
return {"ok": False, "skipped": True, "reason": "disabled"}
|
|
if not settings["base"] or not settings["password"]:
|
|
log(
|
|
"[cpa] remote import enabled but missing cpa_remote_base / "
|
|
"cpa_remote_password (or CPA_REMOTE_BASE / CPA_REMOTE_PASSWORD)"
|
|
)
|
|
return {
|
|
"ok": False,
|
|
"skipped": True,
|
|
"reason": "missing_base_or_password",
|
|
}
|
|
return import_cpa_auth_to_remote(
|
|
path,
|
|
base=settings["base"],
|
|
password=settings["password"],
|
|
retries=settings["retries"],
|
|
delay=settings["delay"],
|
|
log=log,
|
|
)
|
|
|
|
|
|
def export_cpa_xai_for_account(
|
|
email: str,
|
|
password: str,
|
|
*,
|
|
page: Any | None = None,
|
|
cookies: Any | None = None,
|
|
sso: str | None = None,
|
|
config: dict | None = None,
|
|
log_callback: Callable[[str], None] | None = None,
|
|
) -> dict:
|
|
"""Mint OIDC + write xai-<email>.json under register cpa_auths (and optional CPA auth-dir)."""
|
|
cfg = config or {}
|
|
log = log_callback or (lambda m: print(m, flush=True))
|
|
|
|
if not cfg.get("cpa_export_enabled", True):
|
|
log("[cpa] export disabled")
|
|
return {"ok": False, "skipped": True, "reason": "disabled"}
|
|
|
|
tools_dir = cfg.get("api_reverse_tools") or cfg.get("cpa_xai_parent") or None
|
|
_ensure_cpa_xai_on_path(tools_dir)
|
|
|
|
try:
|
|
from cpa_xai import mint_and_export # type: ignore
|
|
except Exception as e: # noqa: BLE001
|
|
log(f"[cpa] import cpa_xai failed: {e}")
|
|
return {"ok": False, "error": f"import: {e}"}
|
|
|
|
out_dir = Path(cfg.get("cpa_auth_dir") or _DEFAULT_OUT).expanduser()
|
|
if not out_dir.is_absolute():
|
|
out_dir = (_REG_DIR / out_dir).resolve()
|
|
|
|
hotload_raw = (cfg.get("cpa_hotload_dir") or "").strip()
|
|
cpa_dir = Path(hotload_raw).expanduser() if hotload_raw else None
|
|
if cpa_dir and not cpa_dir.is_absolute():
|
|
cpa_dir = (_REG_DIR / cpa_dir).resolve()
|
|
|
|
# Priority: cpa_proxy > proxy > env. Config must beat shell https_proxy.
|
|
proxy = (cfg.get("cpa_proxy") or cfg.get("proxy") or "").strip()
|
|
if not proxy:
|
|
proxy = (
|
|
os.environ.get("https_proxy")
|
|
or os.environ.get("HTTPS_PROXY")
|
|
or os.environ.get("http_proxy")
|
|
or ""
|
|
).strip()
|
|
# Default headed when a display exists; headless is frequently Cloudflare-blocked.
|
|
# On servers without DISPLAY, always force headless so Chromium can start.
|
|
has_disp = sys.platform == "win32" or bool(
|
|
(os.environ.get("DISPLAY") or "").strip()
|
|
or (os.environ.get("WAYLAND_DISPLAY") or "").strip()
|
|
)
|
|
cpa_hl_raw = cfg.get("cpa_headless", None)
|
|
if cpa_hl_raw is None:
|
|
# inherit browser_headless auto / explicit if set
|
|
cpa_hl_raw = cfg.get("browser_headless", "auto")
|
|
if isinstance(cpa_hl_raw, bool):
|
|
headless = cpa_hl_raw
|
|
else:
|
|
s = str(cpa_hl_raw or "auto").strip().lower()
|
|
if s in ("1", "true", "yes", "on", "headless"):
|
|
headless = True
|
|
elif s in ("0", "false", "no", "off", "headed", "gui"):
|
|
headless = False
|
|
else:
|
|
headless = not has_disp
|
|
if not has_disp and not headless:
|
|
headless = True
|
|
_log(
|
|
f"cpa browser forced headless (no DISPLAY={os.environ.get('DISPLAY', '')!r}); "
|
|
"prefer xvfb-run for Turnstile; protocol mint still preferred first"
|
|
)
|
|
elif headless and not bool(cfg.get("cpa_headless", False)):
|
|
_log(
|
|
"cpa browser headless auto/on "
|
|
f"(DISPLAY={os.environ.get('DISPLAY', '')!r}); "
|
|
"protocol mint still preferred first"
|
|
)
|
|
probe = bool(cfg.get("cpa_probe_after_write", True))
|
|
probe_chat = bool(cfg.get("cpa_probe_chat", False))
|
|
timeout = float(cfg.get("cpa_mint_timeout_sec", 240))
|
|
base_url = cfg.get("cpa_base_url") or "https://cli-chat-proxy.grok.com/v1"
|
|
force_standalone = bool(cfg.get("cpa_force_standalone", True))
|
|
cookie_inject = bool(cfg.get("cpa_mint_cookie_inject", True))
|
|
reuse_browser = bool(cfg.get("cpa_mint_browser_reuse", True))
|
|
recycle_every = int(cfg.get("cpa_mint_browser_recycle_every", 15) or 0)
|
|
# Protocol (pure HTTP SSO device flow) first; browser only on failure.
|
|
prefer_protocol = bool(cfg.get("cpa_prefer_protocol", True))
|
|
protocol_only = bool(cfg.get("cpa_protocol_only", False))
|
|
protocol_poll_timeout = float(cfg.get("cpa_protocol_poll_timeout_sec", 90) or 90)
|
|
|
|
# cookies: explicit arg > page export > none
|
|
use_cookies = cookies
|
|
if use_cookies is None and cookie_inject and page is not None:
|
|
use_cookies = export_cookies_from_page(page)
|
|
if not cookie_inject:
|
|
use_cookies = None
|
|
else:
|
|
# Always attach SSO cookie clones — register cookies alone often miss accounts.x.ai host
|
|
sso_val = (sso or "").strip()
|
|
if not sso_val and isinstance(use_cookies, list):
|
|
for c in use_cookies:
|
|
if isinstance(c, dict) and c.get("name") in ("sso", "sso-rw") and c.get("value"):
|
|
sso_val = str(c.get("value"))
|
|
break
|
|
if sso_val:
|
|
base = list(use_cookies) if isinstance(use_cookies, list) else []
|
|
for name in ("sso", "sso-rw"):
|
|
for dom in (".x.ai", "accounts.x.ai", ".accounts.x.ai", "auth.x.ai", "grok.com", ".grok.com"):
|
|
base.append({
|
|
"name": name,
|
|
"value": sso_val,
|
|
"domain": dom,
|
|
"path": "/",
|
|
"secure": True,
|
|
"httpOnly": True,
|
|
})
|
|
use_cookies = base
|
|
|
|
sso_val = (sso or "").strip()
|
|
if not sso_val and isinstance(use_cookies, list):
|
|
for c in use_cookies:
|
|
if isinstance(c, dict) and c.get("name") in ("sso", "sso-rw") and c.get("value"):
|
|
sso_val = str(c.get("value"))
|
|
break
|
|
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
log(
|
|
f"[cpa] mint OIDC for {email} -> {out_dir} proxy={proxy or '(none)'} "
|
|
f"cookies={len(use_cookies) if isinstance(use_cookies, list) else (1 if use_cookies else 0)} "
|
|
f"reuse={reuse_browser} protocol={prefer_protocol}"
|
|
f"{' only' if protocol_only else ''} sso={'yes' if sso_val else 'no'}"
|
|
)
|
|
|
|
def _log(msg: str) -> None:
|
|
log(f"[cpa] {msg}")
|
|
|
|
result = mint_and_export(
|
|
email=email,
|
|
password=password,
|
|
auth_dir=out_dir,
|
|
page=None if force_standalone else page,
|
|
proxy=proxy or None,
|
|
headless=headless,
|
|
base_url=base_url,
|
|
probe=probe,
|
|
probe_chat=probe_chat,
|
|
browser_timeout_sec=timeout,
|
|
force_standalone=force_standalone,
|
|
cookies=use_cookies,
|
|
sso=sso_val or None,
|
|
reuse_browser=reuse_browser,
|
|
recycle_every=recycle_every,
|
|
prefer_protocol=prefer_protocol,
|
|
protocol_only=protocol_only,
|
|
protocol_poll_timeout_sec=protocol_poll_timeout,
|
|
log=_log,
|
|
)
|
|
if result.get("mint_method"):
|
|
log(f"[cpa] mint_method={result.get('mint_method')}")
|
|
|
|
# By default, a failed post-write probe is only a warning: the CPA auth file
|
|
# has already been minted and written. Set cpa_probe_required=true to make
|
|
# missing /models grok-4.5 fail the export.
|
|
if (
|
|
not result.get("ok")
|
|
and result.get("path")
|
|
and str(result.get("error") or "").startswith("token ok but grok-4.5 not listed")
|
|
and not cfg.get("cpa_probe_required", False)
|
|
):
|
|
result["ok"] = True
|
|
result["probe_warning"] = result.pop("error", "probe failed")
|
|
log(f"[cpa] probe warning ignored (file already written): {result.get('probe_warning')}")
|
|
|
|
if result.get("ok") and result.get("path") and cfg.get("cpa_copy_to_hotload", False) and cpa_dir:
|
|
try:
|
|
cpa_dir.mkdir(parents=True, exist_ok=True)
|
|
src = Path(result["path"])
|
|
dst = cpa_dir / src.name
|
|
shutil.copy2(src, dst)
|
|
os.chmod(dst, 0o600)
|
|
result["cpa_path"] = str(dst)
|
|
log(f"[cpa] hotload copy -> {dst}")
|
|
except Exception as e: # noqa: BLE001
|
|
log(f"[cpa] hotload copy failed: {e}")
|
|
result["cpa_copy_error"] = str(e)
|
|
|
|
# Optional: push auth file into remote CPA via Management API.
|
|
# Failures never flip result["ok"] — local mint already succeeded.
|
|
if result.get("ok") and result.get("path"):
|
|
try:
|
|
remote = maybe_import_cpa_auth_remote(
|
|
result["path"],
|
|
config=cfg,
|
|
log_callback=log,
|
|
)
|
|
if remote.get("skipped"):
|
|
result["remote_import_skipped"] = remote.get("reason") or "skipped"
|
|
elif remote.get("ok"):
|
|
result["remote_import_ok"] = True
|
|
result["remote_import_base"] = remote.get("base")
|
|
result["remote_import_name"] = remote.get("name")
|
|
else:
|
|
result["remote_import_ok"] = False
|
|
result["remote_import_error"] = remote.get("error") or "unknown"
|
|
except Exception as e: # noqa: BLE001
|
|
log(f"[cpa] remote import unexpected error: {e}")
|
|
result["remote_import_ok"] = False
|
|
result["remote_import_error"] = str(e)
|
|
|
|
# failure log under register dir
|
|
if not result.get("ok"):
|
|
fail_path = out_dir / "cpa_auth_failed.txt"
|
|
with open(fail_path, "a", encoding="utf-8") as f:
|
|
f.write(f"{email}----{result.get('error') or 'unknown'}----{int(time.time())}\n")
|
|
if cfg.get("cpa_mint_required", False):
|
|
raise RuntimeError(f"CPA mint required but failed: {result.get('error')}")
|
|
|
|
return result
|