refactor: modularize registration runtime safely
This commit is contained in:
github-actions[bot]
@@ -0,0 +1,359 @@
|
||||
"""CPA Chromium session lifecycle, reuse, cleanup, and cookies."""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable, Optional
|
||||
|
||||
LogFn = Callable[[str], None]
|
||||
_mint_tls = threading.local()
|
||||
_mint_registry_lock = threading.Lock()
|
||||
_mint_registry = set()
|
||||
|
||||
def _noop_log(_: str) -> None:
|
||||
return None
|
||||
|
||||
class BrowserConfirmError(RuntimeError):
|
||||
pass
|
||||
|
||||
def _sleep(sec: float) -> None:
|
||||
time.sleep(sec)
|
||||
|
||||
def create_standalone_page(proxy: Optional[str] = None, headless: bool = False, log: Optional[LogFn] = None):
|
||||
logger = log or _noop_log
|
||||
try:
|
||||
from DrissionPage import Chromium, ChromiumOptions
|
||||
except ImportError as exc:
|
||||
raise BrowserConfirmError("DrissionPage not installed") from exc
|
||||
|
||||
options = None
|
||||
package_root = Path(__file__).resolve().parents[1]
|
||||
try:
|
||||
register_file = package_root / "grok_register_ttk.py"
|
||||
if register_file.is_file():
|
||||
register_dir = str(package_root)
|
||||
if register_dir not in sys.path:
|
||||
sys.path.insert(0, register_dir)
|
||||
try:
|
||||
from grok_register_ttk import create_browser_options # type: ignore
|
||||
|
||||
options = create_browser_options()
|
||||
logger("using register create_browser_options (turnstilePatch)")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger("register browser options unavailable: %s" % exc)
|
||||
options = None
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger("register options probe failed: %s" % exc)
|
||||
options = None
|
||||
|
||||
if options is None:
|
||||
options = ChromiumOptions()
|
||||
options.auto_port()
|
||||
options.set_timeouts(base=2)
|
||||
for flag in (
|
||||
"--disable-gpu",
|
||||
"--no-sandbox",
|
||||
"--disable-dev-shm-usage",
|
||||
"--mute-audio",
|
||||
"--no-first-run",
|
||||
"--disable-background-networking",
|
||||
"--window-size=1280,900",
|
||||
):
|
||||
options.set_argument(flag)
|
||||
extension = str(package_root / "turnstilePatch")
|
||||
if os.path.isdir(extension):
|
||||
try:
|
||||
options.add_extension(extension)
|
||||
logger("added extension %s" % extension)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
logger("extension add failed: %s" % exc)
|
||||
|
||||
if headless:
|
||||
try:
|
||||
options.headless(True)
|
||||
except Exception:
|
||||
options.set_argument("--headless=new")
|
||||
logger("headless=True (may hit Cloudflare / break real clicks)")
|
||||
else:
|
||||
try:
|
||||
options.headless(False)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
for candidate in (
|
||||
"/usr/bin/chromium",
|
||||
"/usr/bin/chromium-browser",
|
||||
"/usr/bin/google-chrome",
|
||||
"/usr/bin/google-chrome-stable",
|
||||
):
|
||||
if os.path.isfile(candidate):
|
||||
try:
|
||||
options.set_browser_path(candidate)
|
||||
except Exception:
|
||||
pass
|
||||
break
|
||||
|
||||
from .proxyutil import prepare_chromium_proxy, proxy_log_label, resolve_proxy
|
||||
|
||||
resolved = resolve_proxy(proxy)
|
||||
proxy_bridge = None
|
||||
browser = None
|
||||
chrome_proxy, proxy_bridge = prepare_chromium_proxy(resolved, log=logger)
|
||||
try:
|
||||
if chrome_proxy:
|
||||
options.set_argument("--proxy-server=%s" % chrome_proxy)
|
||||
logger("browser proxy=%s (chromium %s)" % (proxy_log_label(resolved), chrome_proxy))
|
||||
else:
|
||||
logger("browser proxy=(none)")
|
||||
browser = Chromium(options)
|
||||
if proxy_bridge is not None:
|
||||
setattr(browser, "_cpa_proxy_bridge", proxy_bridge)
|
||||
page = browser.latest_tab
|
||||
_register_mint_browser(browser)
|
||||
logger("standalone chromium started")
|
||||
return browser, page
|
||||
except Exception:
|
||||
if browser is not None:
|
||||
close_standalone(browser)
|
||||
if proxy_bridge is not None:
|
||||
try:
|
||||
proxy_bridge.stop()
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
|
||||
def close_standalone(browser: Any) -> None:
|
||||
if browser is None:
|
||||
return
|
||||
_unregister_mint_browser(browser)
|
||||
bridge = getattr(browser, "_cpa_proxy_bridge", None)
|
||||
try:
|
||||
browser.quit()
|
||||
except Exception:
|
||||
pass
|
||||
if bridge is not None:
|
||||
try:
|
||||
bridge.stop()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _register_mint_browser(browser: Any) -> None:
|
||||
if browser is None:
|
||||
return
|
||||
with _mint_registry_lock:
|
||||
_mint_registry.add(browser)
|
||||
|
||||
def _unregister_mint_browser(browser: Any) -> None:
|
||||
if browser is None:
|
||||
return
|
||||
with _mint_registry_lock:
|
||||
_mint_registry.discard(browser)
|
||||
|
||||
def _mint_tls_get():
|
||||
state = getattr(_mint_tls, "state", None)
|
||||
if state is None:
|
||||
state = {"browser": None, "page": None, "served": 0, "proxy": None, "headless": None}
|
||||
_mint_tls.state = state
|
||||
return state
|
||||
|
||||
def clear_page_session(page: Any, browser: Optional[Any] = None, log: Optional[LogFn] = None) -> None:
|
||||
logger = log or _noop_log
|
||||
try:
|
||||
if page is not None:
|
||||
try:
|
||||
page.get("about:blank")
|
||||
except Exception:
|
||||
pass
|
||||
for javascript in (
|
||||
"try{localStorage.clear()}catch(e){}",
|
||||
"try{sessionStorage.clear()}catch(e){}",
|
||||
):
|
||||
try:
|
||||
page.run_js(javascript)
|
||||
except Exception:
|
||||
pass
|
||||
for target in (page, browser):
|
||||
if target is None:
|
||||
continue
|
||||
try:
|
||||
target.set.cookies.clear() # type: ignore[attr-defined]
|
||||
logger("mint session cookies cleared")
|
||||
break
|
||||
except Exception:
|
||||
try:
|
||||
cookies = target.cookies()
|
||||
if isinstance(cookies, list):
|
||||
for cookie in cookies:
|
||||
try:
|
||||
target.set.cookies.remove(cookie) # type: ignore[attr-defined]
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
except Exception as exc:
|
||||
logger("clear_page_session: %s" % exc)
|
||||
|
||||
def normalize_cookies(cookies: Any):
|
||||
output = []
|
||||
if not cookies:
|
||||
return output
|
||||
if isinstance(cookies, dict):
|
||||
for name, value in cookies.items():
|
||||
if name and value is not None:
|
||||
output.append({"name": str(name), "value": str(value), "domain": ".x.ai", "path": "/"})
|
||||
cookies = output
|
||||
output = []
|
||||
if not isinstance(cookies, (list, tuple)):
|
||||
return output
|
||||
for cookie in cookies:
|
||||
if not isinstance(cookie, dict):
|
||||
continue
|
||||
name = cookie.get("name") or cookie.get("Name")
|
||||
value = cookie.get("value") or cookie.get("Value")
|
||||
if not name or value is None:
|
||||
continue
|
||||
domain = str(cookie.get("domain") or cookie.get("Domain") or ".x.ai")
|
||||
path = str(cookie.get("path") or cookie.get("Path") or "/")
|
||||
item = {"name": str(name), "value": str(value), "domain": domain, "path": path}
|
||||
for source, target in (
|
||||
("expiry", "expiry"),
|
||||
("expires", "expiry"),
|
||||
("secure", "secure"),
|
||||
("httpOnly", "httpOnly"),
|
||||
("sameSite", "sameSite"),
|
||||
):
|
||||
if source in cookie and cookie[source] is not None:
|
||||
item[target] = cookie[source]
|
||||
output.append(item)
|
||||
sso_names = {"sso", "sso-rw", "cf_clearance", "sso_jwt", "__cf_bm"}
|
||||
extras = []
|
||||
seen = {(item["name"], item["domain"], item["path"]) for item in output}
|
||||
for item in list(output):
|
||||
name = item["name"]
|
||||
if name not in sso_names and not name.startswith("sso"):
|
||||
continue
|
||||
for domain in (".x.ai", "accounts.x.ai", ".accounts.x.ai", "auth.x.ai", ".auth.x.ai"):
|
||||
key = (name, domain, item["path"])
|
||||
if key in seen:
|
||||
continue
|
||||
clone = dict(item)
|
||||
clone["domain"] = domain
|
||||
extras.append(clone)
|
||||
seen.add(key)
|
||||
output.extend(extras)
|
||||
return output
|
||||
|
||||
def inject_cookies(page: Any, cookies: Any, log: Optional[LogFn] = None) -> int:
|
||||
logger = log or _noop_log
|
||||
items = normalize_cookies(cookies)
|
||||
if not items or page is None:
|
||||
return 0
|
||||
for url in ("https://accounts.x.ai/", "https://auth.x.ai/", "https://grok.com/"):
|
||||
try:
|
||||
page.get(url)
|
||||
_sleep(0.4)
|
||||
except Exception:
|
||||
continue
|
||||
count = 0
|
||||
for target_name, target in (("page", page), ("browser", getattr(page, "browser", None))):
|
||||
if target is None:
|
||||
continue
|
||||
try:
|
||||
target.set.cookies(items) # type: ignore[attr-defined]
|
||||
count = len(items)
|
||||
logger("injected cookies bulk via %s=%s" % (target_name, count))
|
||||
break
|
||||
except Exception as exc:
|
||||
logger("bulk set via %s failed: %s" % (target_name, exc))
|
||||
if count == 0:
|
||||
for item in items:
|
||||
ok = False
|
||||
for target in (page, getattr(page, "browser", None)):
|
||||
if target is None:
|
||||
continue
|
||||
try:
|
||||
target.set.cookies(item) # type: ignore[attr-defined]
|
||||
ok = True
|
||||
break
|
||||
except Exception:
|
||||
continue
|
||||
if ok:
|
||||
count += 1
|
||||
logger("injected cookies one-by-one=%s/%s" % (count, len(items)))
|
||||
try:
|
||||
javascript_items = [
|
||||
cookie
|
||||
for cookie in items
|
||||
if not bool(cookie.get("httpOnly")) and str(cookie.get("name") or "").startswith("sso")
|
||||
]
|
||||
if javascript_items:
|
||||
page.run_js(
|
||||
"""
|
||||
const items = arguments[0] || [];
|
||||
for (const c of items) {
|
||||
let cookie = `${c.name}=${c.value}; path=${c.path || '/'}; domain=${c.domain || '.x.ai'}`;
|
||||
if (c.secure !== false) cookie += '; Secure';
|
||||
document.cookie = cookie;
|
||||
}
|
||||
return document.cookie;
|
||||
""",
|
||||
javascript_items,
|
||||
)
|
||||
logger("injected non-httpOnly sso cookies via document.cookie")
|
||||
except Exception as exc:
|
||||
logger("document.cookie injection failed: %s" % exc)
|
||||
return count
|
||||
|
||||
def acquire_mint_browser(proxy: Optional[str] = None, headless: bool = False, reuse: bool = True, recycle_every: int = 15, log: Optional[LogFn] = None):
|
||||
logger = log or _noop_log
|
||||
state = _mint_tls_get()
|
||||
if not reuse:
|
||||
browser, page = create_standalone_page(proxy=proxy, headless=headless, log=logger)
|
||||
return browser, page, True
|
||||
proxy_key = str(proxy or "")
|
||||
if state["browser"] is None:
|
||||
browser, page = create_standalone_page(proxy=proxy, headless=headless, log=logger)
|
||||
state.update({"browser": browser, "page": page, "served": 0, "proxy": proxy_key, "headless": bool(headless)})
|
||||
return browser, page, False
|
||||
if state["proxy"] != proxy_key or bool(state["headless"]) != bool(headless):
|
||||
try:
|
||||
close_standalone(state["browser"])
|
||||
except Exception:
|
||||
pass
|
||||
browser, page = create_standalone_page(proxy=proxy, headless=headless, log=logger)
|
||||
state.update({"browser": browser, "page": page, "served": 0, "proxy": proxy_key, "headless": bool(headless)})
|
||||
return browser, page, False
|
||||
if recycle_every and state["served"] and state["served"] % max(int(recycle_every), 1) == 0:
|
||||
try:
|
||||
close_standalone(state["browser"])
|
||||
except Exception:
|
||||
pass
|
||||
browser, page = create_standalone_page(proxy=proxy, headless=headless, log=logger)
|
||||
state.update({"browser": browser, "page": page, "served": 0, "proxy": proxy_key, "headless": bool(headless)})
|
||||
return browser, page, False
|
||||
clear_page_session(state["page"], browser=state["browser"], log=logger)
|
||||
return state["browser"], state["page"], False
|
||||
|
||||
def release_mint_browser(owned: bool, success: bool, log: Optional[LogFn] = None) -> None:
|
||||
logger = log or _noop_log
|
||||
if owned:
|
||||
return
|
||||
state = _mint_tls_get()
|
||||
if success:
|
||||
state["served"] = int(state.get("served", 0) or 0) + 1
|
||||
logger("mint browser served=%s" % state["served"])
|
||||
|
||||
def shutdown_mint_browsers() -> None:
|
||||
state = _mint_tls_get()
|
||||
with _mint_registry_lock:
|
||||
browsers = list(_mint_registry)
|
||||
for browser in browsers:
|
||||
try:
|
||||
close_standalone(browser)
|
||||
except Exception:
|
||||
pass
|
||||
state.update({"browser": None, "page": None, "served": 0, "proxy": None, "headless": None})
|
||||
|
||||
Reference in New Issue
Block a user