291 lines
10 KiB
Python
291 lines
10 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""Auto-extract cookies from local browsers for all supported platforms.
|
||
|
||
Supports: Chrome, Firefox, Edge, Brave, Opera
|
||
Extracts: Twitter, XiaoHongShu, Bilibili cookies in one shot.
|
||
|
||
Usage:
|
||
agent-reach configure --from-browser chrome
|
||
"""
|
||
|
||
import sys
|
||
from typing import Dict, List, Optional, Tuple
|
||
|
||
|
||
# Platform cookie specs: (platform_name, domain_pattern, needed_cookies)
|
||
PLATFORM_SPECS = [
|
||
{
|
||
"name": "Twitter/X",
|
||
"domains": [".x.com", ".twitter.com"],
|
||
"cookies": ["auth_token", "ct0"],
|
||
"config_key": "twitter",
|
||
},
|
||
{
|
||
"name": "XiaoHongShu",
|
||
"domains": [".xiaohongshu.com"],
|
||
"cookies": None, # None = grab all cookies as header string
|
||
"config_key": "xhs",
|
||
},
|
||
{
|
||
"name": "Bilibili",
|
||
"domains": [".bilibili.com"],
|
||
"cookies": ["SESSDATA", "bili_jct"],
|
||
"config_key": "bilibili",
|
||
},
|
||
{
|
||
"name": "Xueqiu",
|
||
"domains": [".xueqiu.com", "xueqiu.com"],
|
||
"cookies": None, # grab all — xq_a_token + session cookies required
|
||
"config_key": "xueqiu",
|
||
},
|
||
]
|
||
|
||
|
||
def extract_all(browser: str = "chrome") -> Dict[str, dict]:
|
||
"""
|
||
Extract cookies for all supported platforms from the specified browser.
|
||
|
||
Returns:
|
||
{
|
||
"twitter": {"auth_token": "xxx", "ct0": "yyy"},
|
||
"xhs": {"cookie_string": "a=1; b=2; ..."},
|
||
"bilibili": {"SESSDATA": "xxx", "bili_jct": "yyy"},
|
||
}
|
||
"""
|
||
# Try rookiepy first (Rust-based, more stable), fallback to browser_cookie3
|
||
use_rookiepy = False
|
||
try:
|
||
import rookiepy
|
||
use_rookiepy = True
|
||
except ImportError:
|
||
try:
|
||
import browser_cookie3
|
||
except ImportError:
|
||
raise RuntimeError(
|
||
"Cookie extraction requires rookiepy or browser_cookie3.\n"
|
||
"Install: pip install rookiepy (recommended)\n"
|
||
" or: pip install browser-cookie3"
|
||
)
|
||
|
||
browser = browser.lower()
|
||
supported = ["chrome", "firefox", "edge", "brave", "opera"]
|
||
if browser not in supported:
|
||
raise ValueError(
|
||
f"Unsupported browser: {browser}. Supported: {', '.join(supported)}"
|
||
)
|
||
|
||
if use_rookiepy:
|
||
# rookiepy returns list of dicts with name/value/domain/path keys
|
||
try:
|
||
browser_funcs = {
|
||
"chrome": rookiepy.chrome,
|
||
"firefox": rookiepy.firefox,
|
||
"edge": rookiepy.edge,
|
||
"brave": rookiepy.brave,
|
||
"opera": rookiepy.opera,
|
||
}
|
||
raw_cookies = browser_funcs[browser]()
|
||
# Wrap into objects with .name, .value, .domain for compatibility
|
||
class _Cookie:
|
||
def __init__(self, d):
|
||
self.name = d.get("name", "")
|
||
self.value = d.get("value", "")
|
||
self.domain = d.get("domain", "")
|
||
cookie_jar = [_Cookie(c) for c in raw_cookies]
|
||
except Exception as e:
|
||
raise RuntimeError(
|
||
f"Could not read {browser} cookies via rookiepy: {e}\n"
|
||
f"Make sure {browser} is closed and you have permission."
|
||
)
|
||
else:
|
||
browser_funcs = {
|
||
"chrome": browser_cookie3.chrome,
|
||
"firefox": browser_cookie3.firefox,
|
||
"edge": browser_cookie3.edge,
|
||
"brave": browser_cookie3.brave,
|
||
"opera": browser_cookie3.opera,
|
||
}
|
||
try:
|
||
cookie_jar = browser_funcs[browser]()
|
||
except Exception as e:
|
||
raise RuntimeError(
|
||
f"Could not read {browser} cookies: {e}\n"
|
||
f"Make sure {browser} is closed and you have permission."
|
||
)
|
||
|
||
results = {}
|
||
|
||
for spec in PLATFORM_SPECS:
|
||
platform_cookies = {}
|
||
all_cookies_for_domain = []
|
||
|
||
for cookie in cookie_jar:
|
||
# Check if cookie belongs to this platform
|
||
domain_match = any(
|
||
cookie.domain.endswith(d) or cookie.domain == d.lstrip(".")
|
||
for d in spec["domains"]
|
||
)
|
||
if not domain_match:
|
||
continue
|
||
|
||
all_cookies_for_domain.append(cookie)
|
||
|
||
if spec["cookies"] is not None:
|
||
if cookie.name in spec["cookies"]:
|
||
platform_cookies[cookie.name] = cookie.value
|
||
|
||
if spec["cookies"] is None:
|
||
# Grab all as header string
|
||
if all_cookies_for_domain:
|
||
cookie_str = "; ".join(
|
||
f"{c.name}={c.value}" for c in all_cookies_for_domain
|
||
)
|
||
results[spec["config_key"]] = {"cookie_string": cookie_str}
|
||
else:
|
||
if platform_cookies:
|
||
results[spec["config_key"]] = platform_cookies
|
||
|
||
return results
|
||
|
||
|
||
def _open_owner_only(path: str):
|
||
"""Open *path* for writing, atomically creating it with mode 0o600.
|
||
|
||
Mirrors the pattern used by Config.save() in config.py: O_WRONLY|O_CREAT|
|
||
O_TRUNC + an explicit mode argument so the file is never briefly
|
||
world-readable between open() and a later os.chmod(). On Windows (or any
|
||
OS that rejects the open flags) we fall back to a plain open().
|
||
"""
|
||
import os
|
||
import stat
|
||
|
||
try:
|
||
fd = os.open(
|
||
path,
|
||
os.O_WRONLY | os.O_CREAT | os.O_TRUNC,
|
||
stat.S_IRUSR | stat.S_IWUSR, # 0o600
|
||
)
|
||
return os.fdopen(fd, "w", encoding="utf-8")
|
||
except OSError:
|
||
return open(path, "w", encoding="utf-8")
|
||
|
||
|
||
def _sync_xfetch_session(auth_token: str, ct0: str) -> None:
|
||
"""Sync Twitter credentials to ~/.config/xfetch/session.json (legacy xreach compat)."""
|
||
import json
|
||
import os
|
||
|
||
try:
|
||
xfetch_dir = os.path.join(os.path.expanduser("~"), ".config", "xfetch")
|
||
os.makedirs(xfetch_dir, exist_ok=True)
|
||
session_path = os.path.join(xfetch_dir, "session.json")
|
||
session_data: dict = {}
|
||
if os.path.exists(session_path):
|
||
try:
|
||
with open(session_path, "r", encoding="utf-8") as sf:
|
||
session_data = json.load(sf)
|
||
except (json.JSONDecodeError, OSError):
|
||
session_data = {}
|
||
session_data["authToken"] = auth_token
|
||
session_data["ct0"] = ct0
|
||
with _open_owner_only(session_path) as sf:
|
||
json.dump(session_data, sf, indent=2)
|
||
except Exception:
|
||
# Non-fatal: agent-reach config is the source of truth, xfetch sync is best-effort
|
||
pass
|
||
|
||
|
||
def _sync_bird_env(auth_token: str, ct0: str) -> None:
|
||
"""Write Twitter credentials to ~/.config/bird/credentials.env for bird CLI.
|
||
|
||
bird reads AUTH_TOKEN and CT0 from environment variables. This writes a
|
||
shell-sourceable file so users can `source ~/.config/bird/credentials.env`.
|
||
Values are passed through shlex.quote so a token containing a quote, $, or
|
||
backtick cannot break out into shell syntax when the file is sourced.
|
||
"""
|
||
import os
|
||
import shlex
|
||
|
||
try:
|
||
bird_dir = os.path.join(os.path.expanduser("~"), ".config", "bird")
|
||
os.makedirs(bird_dir, exist_ok=True)
|
||
env_path = os.path.join(bird_dir, "credentials.env")
|
||
with _open_owner_only(env_path) as f:
|
||
f.write(f"AUTH_TOKEN={shlex.quote(auth_token)}\n")
|
||
f.write(f"CT0={shlex.quote(ct0)}\n")
|
||
except Exception:
|
||
# Non-fatal: agent-reach config is the source of truth, bird env sync is best-effort
|
||
pass
|
||
|
||
|
||
# Alias for callers expecting the name _sync_bird_credentials
|
||
_sync_bird_credentials = _sync_bird_env
|
||
|
||
|
||
def configure_from_browser(browser: str, config) -> List[Tuple[str, bool, str]]:
|
||
"""
|
||
Extract cookies and configure all found platforms.
|
||
|
||
Returns list of (platform_name, success, message) tuples.
|
||
"""
|
||
results_list = []
|
||
|
||
try:
|
||
extracted = extract_all(browser)
|
||
except Exception as e:
|
||
return [("Browser", False, str(e))]
|
||
|
||
if not extracted:
|
||
return [("All platforms", False,
|
||
f"No platform cookies found in {browser}. "
|
||
f"Make sure you're logged into Twitter, XiaoHongShu, etc. in {browser}.")]
|
||
|
||
# Configure each found platform
|
||
if "twitter" in extracted:
|
||
tc = extracted["twitter"]
|
||
if "auth_token" in tc and "ct0" in tc:
|
||
config.set("twitter_auth_token", tc["auth_token"])
|
||
config.set("twitter_ct0", tc["ct0"])
|
||
# Legacy sync (best-effort)
|
||
_sync_xfetch_session(tc["auth_token"], tc["ct0"])
|
||
results_list.append(("Twitter/X", True, "auth_token + ct0"))
|
||
else:
|
||
found = ", ".join(tc.keys())
|
||
missing = [k for k in ["auth_token", "ct0"] if k not in tc]
|
||
results_list.append(("Twitter/X", False,
|
||
f"Found {found}, but missing: {', '.join(missing)}. "
|
||
f"Make sure you're logged into x.com in {browser}."))
|
||
|
||
if "xhs" in extracted:
|
||
cookie_str = extracted["xhs"].get("cookie_string", "")
|
||
if cookie_str:
|
||
config.set("xhs_cookie", cookie_str)
|
||
n_cookies = len(cookie_str.split(";"))
|
||
results_list.append(("XiaoHongShu", True, f"{n_cookies} cookies"))
|
||
|
||
if "bilibili" in extracted:
|
||
bc = extracted["bilibili"]
|
||
if "SESSDATA" in bc:
|
||
config.set("bilibili_sessdata", bc["SESSDATA"])
|
||
if "bili_jct" in bc:
|
||
config.set("bilibili_csrf", bc["bili_jct"])
|
||
results_list.append(("Bilibili", True, "SESSDATA" +
|
||
(" + bili_jct" if "bili_jct" in bc else "")))
|
||
else:
|
||
results_list.append(("Bilibili", False,
|
||
f"No SESSDATA found. Make sure you're logged into bilibili.com in {browser}."))
|
||
|
||
if "xueqiu" in extracted:
|
||
cookie_str = extracted["xueqiu"].get("cookie_string", "")
|
||
# Only save if xq_a_token is present — anonymous cookies are useless
|
||
if cookie_str and "xq_a_token" in cookie_str:
|
||
config.set("xueqiu_cookie", cookie_str)
|
||
n_cookies = len(cookie_str.split(";"))
|
||
results_list.append(("Xueqiu", True, f"{n_cookies} cookies (含 xq_a_token)"))
|
||
elif cookie_str:
|
||
results_list.append(("Xueqiu", False,
|
||
f"找到 {len(cookie_str.split(';'))} 个 Cookie 但缺少 xq_a_token，"
|
||
f"请先在 {browser} 中登录 xueqiu.com"))
|
||
|
||
return results_list
|