Ubuntu
527 lines
16 KiB
Python
527 lines
16 KiB
Python
"""Cookie extraction and management for XHS API client."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import functools
|
|
import json
|
|
import logging
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
from collections import OrderedDict
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from .constants import CONFIG_DIR_NAME, COOKIE_FILE, INDEX_CACHE_FILE, TOKEN_CACHE_FILE
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Cookie TTL: warn and attempt browser refresh after 7 days
|
|
COOKIE_TTL_DAYS = 7
|
|
_COOKIE_TTL_SECONDS = COOKIE_TTL_DAYS * 86400
|
|
_TOKEN_CACHE_LOCK = threading.RLock()
|
|
_TOKEN_CACHE_MEMORY: OrderedDict[str, dict[str, Any]] | None = None
|
|
_TOKEN_CACHE_PATH: Path | None = None
|
|
NOTE_CONTEXT_TTL_SECONDS = 86400
|
|
|
|
|
|
|
|
def get_config_dir() -> Path:
|
|
"""Get or create config directory."""
|
|
config_dir = Path.home() / CONFIG_DIR_NAME
|
|
config_dir.mkdir(parents=True, exist_ok=True)
|
|
return config_dir
|
|
|
|
|
|
def get_cookie_path() -> Path:
|
|
"""Get cookie file path."""
|
|
return get_config_dir() / COOKIE_FILE
|
|
|
|
|
|
def get_token_cache_path() -> Path:
|
|
"""Get xsec token cache file path."""
|
|
return get_config_dir() / TOKEN_CACHE_FILE
|
|
|
|
|
|
def get_index_cache_path() -> Path:
|
|
"""Get note index cache file path."""
|
|
return get_config_dir() / INDEX_CACHE_FILE
|
|
|
|
|
|
def load_saved_cookies() -> dict[str, str] | None:
|
|
"""Load cookies from local storage."""
|
|
cookie_path = get_cookie_path()
|
|
if not cookie_path.exists():
|
|
return None
|
|
try:
|
|
data = json.loads(cookie_path.read_text())
|
|
if data.get("a1"):
|
|
logger.debug("Loaded saved cookies from %s", cookie_path)
|
|
return data
|
|
except (OSError, json.JSONDecodeError) as e:
|
|
logger.debug("Failed to load saved cookies: %s", e)
|
|
return None
|
|
|
|
|
|
def save_cookies(cookies: dict[str, str]) -> None:
|
|
"""Save cookies to local storage with restricted permissions and TTL timestamp."""
|
|
cookie_path = get_cookie_path()
|
|
payload = {**cookies, "saved_at": time.time()}
|
|
cookie_path.write_text(json.dumps(payload, indent=2))
|
|
cookie_path.chmod(0o600)
|
|
logger.debug("Saved cookies to %s", cookie_path)
|
|
|
|
|
|
def clear_cookies() -> None:
|
|
"""Remove saved cookies."""
|
|
cookie_path = get_cookie_path()
|
|
if cookie_path.exists():
|
|
cookie_path.unlink()
|
|
logger.debug("Cleared cookies from %s", cookie_path)
|
|
|
|
|
|
def _normalize_token_entry(value: Any) -> dict[str, Any] | None:
|
|
if isinstance(value, str):
|
|
return {"token": value, "source": "", "ts": time.time()}
|
|
if not isinstance(value, dict):
|
|
return None
|
|
|
|
token = str(value.get("token", "")).strip()
|
|
if not token:
|
|
return None
|
|
|
|
source = str(value.get("source", "")).strip()
|
|
context = str(value.get("context", "")).strip()
|
|
ts = value.get("ts", 0)
|
|
try:
|
|
ts = float(ts)
|
|
except (TypeError, ValueError):
|
|
ts = 0.0
|
|
|
|
entry = {"token": token, "source": source, "ts": ts}
|
|
if context:
|
|
entry["context"] = context
|
|
return entry
|
|
|
|
|
|
def _load_token_cache_from_disk(cache_path: Path) -> OrderedDict[str, dict[str, Any]]:
|
|
if not cache_path.exists():
|
|
return OrderedDict()
|
|
try:
|
|
data = json.loads(cache_path.read_text())
|
|
except (OSError, json.JSONDecodeError) as exc:
|
|
logger.debug("Failed to load token cache: %s", exc)
|
|
return OrderedDict()
|
|
if not isinstance(data, dict):
|
|
return OrderedDict()
|
|
|
|
normalized: list[tuple[str, dict[str, Any]]] = []
|
|
for key, value in data.items():
|
|
if not key:
|
|
continue
|
|
entry = _normalize_token_entry(value)
|
|
if entry:
|
|
normalized.append((str(key), entry))
|
|
normalized.sort(key=lambda item: float(item[1].get("ts", 0)))
|
|
return OrderedDict(normalized)
|
|
|
|
|
|
def _prune_token_cache(
|
|
cache: OrderedDict[str, dict[str, Any]],
|
|
now: float | None = None,
|
|
) -> OrderedDict[str, dict[str, Any]]:
|
|
now = now or time.time()
|
|
pruned = OrderedDict(
|
|
(key, value)
|
|
for key, value in cache.items()
|
|
if now - float(value.get("ts", 0)) <= NOTE_CONTEXT_TTL_SECONDS
|
|
)
|
|
while len(pruned) > TOKEN_CACHE_MAX_SIZE:
|
|
pruned.popitem(last=False)
|
|
return pruned
|
|
|
|
|
|
def load_token_cache() -> dict[str, dict[str, Any]]:
|
|
"""Load cached note_id -> token context mappings."""
|
|
cache_path = get_token_cache_path()
|
|
global _TOKEN_CACHE_MEMORY, _TOKEN_CACHE_PATH
|
|
|
|
with _TOKEN_CACHE_LOCK:
|
|
if _TOKEN_CACHE_MEMORY is None or _TOKEN_CACHE_PATH != cache_path:
|
|
_TOKEN_CACHE_MEMORY = _prune_token_cache(_load_token_cache_from_disk(cache_path))
|
|
_TOKEN_CACHE_PATH = cache_path
|
|
return {
|
|
key: dict(value)
|
|
for key, value in _TOKEN_CACHE_MEMORY.items()
|
|
}
|
|
|
|
|
|
def save_token_cache(cache: dict[str, dict[str, Any]]) -> None:
|
|
"""Persist xsec token cache with restricted permissions."""
|
|
cache_path = get_token_cache_path()
|
|
global _TOKEN_CACHE_MEMORY, _TOKEN_CACHE_PATH
|
|
|
|
normalized = _prune_token_cache(OrderedDict(
|
|
sorted(
|
|
(
|
|
(str(key), dict(value))
|
|
for key, value in cache.items()
|
|
if key and isinstance(value, dict)
|
|
),
|
|
key=lambda item: float(item[1].get("ts", 0)),
|
|
)
|
|
))
|
|
|
|
with _TOKEN_CACHE_LOCK:
|
|
cache_path.write_text(json.dumps(normalized, indent=2))
|
|
cache_path.chmod(0o600)
|
|
_TOKEN_CACHE_MEMORY = normalized
|
|
_TOKEN_CACHE_PATH = cache_path
|
|
|
|
|
|
TOKEN_CACHE_MAX_SIZE = 500
|
|
|
|
|
|
def cache_note_context(
|
|
note_id: str,
|
|
xsec_token: str,
|
|
xsec_source: str = "",
|
|
*,
|
|
context: str = "",
|
|
) -> None:
|
|
"""Store a resolved note token and source for later access.
|
|
|
|
Maintains an LRU-style cache capped at TOKEN_CACHE_MAX_SIZE entries.
|
|
Each entry stores token/source/timestamp metadata; overflow evicts the
|
|
oldest entries.
|
|
"""
|
|
if not note_id or not xsec_token:
|
|
return
|
|
cache = load_token_cache()
|
|
|
|
existing = cache.get(note_id)
|
|
if (
|
|
isinstance(existing, dict)
|
|
and existing.get("token") == xsec_token
|
|
and existing.get("source", "") == xsec_source
|
|
and existing.get("context", "") == context
|
|
):
|
|
existing["ts"] = time.time()
|
|
save_token_cache(cache)
|
|
return
|
|
|
|
entry = {
|
|
"token": xsec_token,
|
|
"source": xsec_source,
|
|
"ts": time.time(),
|
|
}
|
|
if context:
|
|
entry["context"] = context
|
|
cache[note_id] = entry
|
|
|
|
# Evict oldest entries if over limit
|
|
if len(cache) > TOKEN_CACHE_MAX_SIZE:
|
|
sorted_keys = sorted(
|
|
cache.keys(),
|
|
key=lambda k: cache[k].get("ts", 0) if isinstance(cache[k], dict) else 0,
|
|
)
|
|
for key in sorted_keys[: len(cache) - TOKEN_CACHE_MAX_SIZE]:
|
|
del cache[key]
|
|
|
|
save_token_cache(cache)
|
|
logger.debug("Cached xsec_token for note %s", note_id)
|
|
|
|
|
|
def invalidate_note_context(note_id: str) -> None:
|
|
"""Remove cached token/source metadata for a note ID."""
|
|
if not note_id:
|
|
return
|
|
cache = load_token_cache()
|
|
if note_id not in cache:
|
|
return
|
|
del cache[note_id]
|
|
save_token_cache(cache)
|
|
logger.debug("Invalidated cached note context for %s", note_id)
|
|
|
|
|
|
def _normalize_index_entry(value: Any) -> dict[str, str] | None:
|
|
if not isinstance(value, dict):
|
|
return None
|
|
|
|
note_id = str(value.get("note_id", "")).strip()
|
|
if not note_id:
|
|
return None
|
|
|
|
return {
|
|
"note_id": note_id,
|
|
"xsec_token": str(value.get("xsec_token", "")).strip(),
|
|
"xsec_source": str(value.get("xsec_source", "")).strip(),
|
|
}
|
|
|
|
|
|
def save_note_index(items: list[dict[str, str]]) -> None:
|
|
"""Persist the latest ordered note index for short-index navigation."""
|
|
path = get_index_cache_path()
|
|
normalized = [
|
|
entry
|
|
for entry in (_normalize_index_entry(item) for item in items)
|
|
if entry
|
|
]
|
|
path.write_text(json.dumps(normalized, indent=2, ensure_ascii=False))
|
|
path.chmod(0o600)
|
|
logger.debug("Saved note index with %d entries", len(normalized))
|
|
|
|
|
|
def get_note_by_index(index: int) -> dict[str, str] | None:
|
|
"""Resolve a 1-based short index to a cached note reference."""
|
|
if index <= 0:
|
|
return None
|
|
|
|
path = get_index_cache_path()
|
|
if not path.exists():
|
|
return None
|
|
|
|
try:
|
|
data = json.loads(path.read_text())
|
|
except (OSError, json.JSONDecodeError):
|
|
return None
|
|
|
|
if not isinstance(data, list) or index > len(data):
|
|
return None
|
|
|
|
return _normalize_index_entry(data[index - 1])
|
|
|
|
|
|
def cache_xsec_token(note_id: str, xsec_token: str) -> None:
|
|
"""Backwards-compatible wrapper for token-only caching."""
|
|
cache_note_context(note_id, xsec_token)
|
|
|
|
|
|
def get_cached_note_context(note_id: str) -> dict[str, Any]:
|
|
"""Get cached token/source metadata for a note ID."""
|
|
entry = load_token_cache().get(note_id)
|
|
if not isinstance(entry, dict):
|
|
return {}
|
|
return {
|
|
"token": str(entry.get("token", "")),
|
|
"source": str(entry.get("source", "")),
|
|
"context": str(entry.get("context", "")),
|
|
"ts": entry.get("ts", 0.0),
|
|
}
|
|
|
|
|
|
def get_cached_xsec_token(note_id: str) -> str:
|
|
"""Get a cached xsec token for a note ID."""
|
|
return get_cached_note_context(note_id).get("token", "")
|
|
|
|
|
|
@functools.lru_cache(maxsize=1)
|
|
def _available_browsers() -> tuple[str, ...]:
|
|
"""List all browser names supported by browser_cookie3 (cached)."""
|
|
import inspect
|
|
|
|
import browser_cookie3 as bc3
|
|
|
|
return tuple(sorted(
|
|
name
|
|
for name in dir(bc3)
|
|
if not name.startswith("_")
|
|
and name != "load" # 'load' tries all browsers; skip in auto-detect
|
|
and callable(getattr(bc3, name))
|
|
and hasattr(getattr(bc3, name), "__code__")
|
|
and "domain_name" in inspect.signature(getattr(bc3, name)).parameters
|
|
))
|
|
|
|
|
|
def _get_browser_loader(source: str):
|
|
"""Get browser cookie loader from browser_cookie3 by name."""
|
|
import browser_cookie3 as bc3
|
|
|
|
loader = getattr(bc3, source, None)
|
|
if loader is None or not callable(loader):
|
|
available = _available_browsers()
|
|
raise ValueError(
|
|
f"Unknown browser: {source!r}. Available: {', '.join(available)}"
|
|
)
|
|
return loader
|
|
|
|
|
|
def _extract_in_process(source: str) -> dict[str, str] | None:
|
|
"""Extract cookies in-process for macOS Keychain compatibility."""
|
|
try:
|
|
loader = _get_browser_loader(source)
|
|
except ImportError:
|
|
logger.debug("browser_cookie3 not installed, skipping in-process extraction")
|
|
return None
|
|
except ValueError as exc:
|
|
logger.debug("%s", exc)
|
|
return None
|
|
|
|
try:
|
|
jar = loader(domain_name=".xiaohongshu.com")
|
|
except Exception as exc:
|
|
logger.debug("%s in-process extraction failed: %s", source, exc)
|
|
return None
|
|
|
|
cookies = {cookie.name: cookie.value for cookie in jar if "xiaohongshu.com" in (cookie.domain or "")}
|
|
if cookies.get("a1"):
|
|
logger.debug("Loaded XHS cookies from %s in-process", source)
|
|
return cookies
|
|
|
|
logger.debug("No usable a1 cookie found in %s in-process extraction", source)
|
|
return None
|
|
|
|
|
|
def _extract_via_subprocess(source: str) -> dict[str, str] | None:
|
|
"""Extract cookies via subprocess to avoid browser SQLite locks."""
|
|
extract_script = '''
|
|
import json, sys
|
|
try:
|
|
import browser_cookie3 as bc3
|
|
except ImportError:
|
|
print(json.dumps({"error": "browser-cookie3 not installed"}))
|
|
sys.exit(0)
|
|
|
|
source = sys.argv[1]
|
|
loader = getattr(bc3, source, None)
|
|
if not loader or not callable(loader):
|
|
print(json.dumps({"error": f"Unknown browser: {source}"}))
|
|
sys.exit(0)
|
|
|
|
try:
|
|
cj = loader(domain_name=".xiaohongshu.com")
|
|
cookies = {c.name: c.value for c in cj if "xiaohongshu.com" in (c.domain or "")}
|
|
if cookies.get("a1"):
|
|
print(json.dumps({"browser": source, "cookies": cookies}))
|
|
else:
|
|
print(json.dumps({"error": "no_a1_cookie"}))
|
|
except Exception as e:
|
|
print(json.dumps({"error": str(e)}))
|
|
'''
|
|
|
|
try:
|
|
result = subprocess.run(
|
|
[sys.executable, "-c", extract_script, source],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=15,
|
|
)
|
|
|
|
if result.returncode != 0:
|
|
logger.debug("Cookie extraction subprocess failed: %s", result.stderr)
|
|
return None
|
|
|
|
data = json.loads(result.stdout.strip())
|
|
if "error" in data:
|
|
logger.debug("Cookie extraction error: %s", data["error"])
|
|
return None
|
|
|
|
return data["cookies"]
|
|
|
|
except subprocess.TimeoutExpired:
|
|
logger.debug("Cookie extraction timed out")
|
|
return None
|
|
except (json.JSONDecodeError, KeyError) as e:
|
|
logger.debug("Cookie extraction parse error: %s", e)
|
|
return None
|
|
|
|
|
|
def extract_browser_cookies(source: str = "auto") -> tuple[str, dict[str, str]] | None:
|
|
"""
|
|
Extract XHS cookies from browser using browser-cookie3.
|
|
|
|
When *source* is ``"auto"``, tries supported browsers with a small
|
|
thread pool and returns the first one that has valid cookies.
|
|
|
|
Returns ``(browser_name, cookies)`` on success, or ``None``.
|
|
"""
|
|
if source != "auto":
|
|
cookies = _extract_in_process(source)
|
|
if cookies:
|
|
return source, cookies
|
|
cookies = _extract_via_subprocess(source)
|
|
if cookies:
|
|
return source, cookies
|
|
return None
|
|
|
|
# Auto-detect: try all available browsers
|
|
try:
|
|
browsers = _available_browsers()
|
|
except ImportError:
|
|
logger.debug("browser_cookie3 not installed")
|
|
return None
|
|
|
|
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
|
|
|
|
def _try_browser(browser: str) -> tuple[str, dict[str, str]] | None:
|
|
logger.debug("Auto-detect: trying %s …", browser)
|
|
cookies = _extract_in_process(browser)
|
|
if cookies:
|
|
return browser, cookies
|
|
cookies = _extract_via_subprocess(browser)
|
|
if cookies:
|
|
return browser, cookies
|
|
return None
|
|
|
|
with ThreadPoolExecutor(max_workers=min(4, len(browsers) or 1)) as pool:
|
|
pending = {pool.submit(_try_browser, browser) for browser in browsers}
|
|
while pending:
|
|
done, pending = wait(pending, return_when=FIRST_COMPLETED)
|
|
for future in done:
|
|
result = future.result()
|
|
if result:
|
|
for rest in pending:
|
|
rest.cancel()
|
|
return result
|
|
|
|
return None
|
|
|
|
|
|
def get_cookies(
|
|
cookie_source: str = "auto", *, force_refresh: bool = False
|
|
) -> tuple[str, dict[str, str]]:
|
|
"""
|
|
Multi-strategy cookie acquisition with TTL-based auto-refresh.
|
|
|
|
Returns ``(browser_name, cookies)``.
|
|
|
|
1. Load saved cookies (skip if stale > 7 days)
|
|
2. Extract from browser (auto-detect if *cookie_source* is ``"auto"``)
|
|
3. Raise error if all fail
|
|
"""
|
|
# 1. Try saved cookies first
|
|
if not force_refresh:
|
|
saved = load_saved_cookies()
|
|
if saved:
|
|
saved_at = saved.pop("saved_at", 0)
|
|
if saved_at and (time.time() - float(saved_at)) > _COOKIE_TTL_SECONDS:
|
|
logger.info(
|
|
"Cookies older than %d days, attempting browser refresh",
|
|
COOKIE_TTL_DAYS,
|
|
)
|
|
result = extract_browser_cookies(cookie_source)
|
|
if result:
|
|
save_cookies(result[1])
|
|
return result
|
|
logger.warning(
|
|
"Cookie refresh failed; using existing cookies (age: %d+ days)",
|
|
COOKIE_TTL_DAYS,
|
|
)
|
|
return "saved", saved
|
|
|
|
# 2. Try browser extraction
|
|
from .exceptions import NoCookieError
|
|
|
|
result = extract_browser_cookies(cookie_source)
|
|
if result:
|
|
save_cookies(result[1])
|
|
return result
|
|
|
|
raise NoCookieError(cookie_source)
|
|
|
|
|
|
def cookies_to_string(cookies: dict[str, str]) -> str:
|
|
"""Format cookies as a cookie header string."""
|
|
return "; ".join(f"{k}={v}" for k, v in cookies.items())
|