diff --git a/tests/test_cookies.py b/tests/test_cookies.py index aae7636..53ecf93 100644 --- a/tests/test_cookies.py +++ b/tests/test_cookies.py @@ -1,15 +1,22 @@ """Unit tests for cookie management (no network required).""" +import sys import time +import types +from pathlib import Path import pytest from xhs_cli.cookies import ( NOTE_CONTEXT_TTL_SECONDS, + _available_browsers, + _cookies_for_domain, + _extract_via_subprocess, cache_note_context, clear_cookies, cookies_to_string, + extract_browser_cookies, get_cached_note_context, get_cached_xsec_token, get_cookies, @@ -84,6 +91,26 @@ def test_format(self): assert "; " in result +class TestCookieDomainFiltering: + def test_matches_exact_domain_and_subdomains_only(self): + class FakeCookie: + def __init__(self, name: str, value: str, domain: str): + self.name = name + self.value = value + self.domain = domain + + cookies = _cookies_for_domain( + [ + FakeCookie("exact", "1", ".xiaohongshu.com"), + FakeCookie("subdomain", "2", ".www.xiaohongshu.com"), + FakeCookie("unrelated", "3", ".fake-xiaohongshu.com"), + ], + ".xiaohongshu.com", + ) + + assert cookies == {"exact": "1", "subdomain": "2"} + + class TestGetCookies: def test_prefers_saved_cookies_by_default(self, monkeypatch): monkeypatch.setattr("xhs_cli.cookies.load_saved_cookies", lambda: {"a1": "saved"}) @@ -110,6 +137,96 @@ def test_force_refresh_bypasses_saved_cookies(self, monkeypatch): assert cookies == {"a1": "fresh"} assert saved == [{"a1": "fresh"}] + def test_extract_browser_cookies_supports_comet(self, tmp_path, monkeypatch): + comet_dir = tmp_path / "Library" / "Application Support" / "Comet" + default_dir = comet_dir / "Default" + default_dir.mkdir(parents=True) + (default_dir / "Cookies").write_text("placeholder") + + captured: list[dict[str, object]] = [] + + class FakeCookie: + def __init__(self, name: str, value: str, domain: str): + self.name = name + self.value = value + self.domain = domain + + class FakeChromiumBased: + def __init__(self, **kwargs): + captured.append(kwargs) + + def load(self): + return [ + FakeCookie("a1", "comet-a1", ".xiaohongshu.com"), + FakeCookie("web_session", "session", ".xiaohongshu.com"), + ] + + fake_bc3 = types.SimpleNamespace(ChromiumBased=FakeChromiumBased) + monkeypatch.setitem(sys.modules, "browser_cookie3", fake_bc3) + monkeypatch.setattr(Path, "home", lambda: tmp_path) + + browser, cookies = extract_browser_cookies("comet") + + assert browser == "comet" + assert cookies["a1"] == "comet-a1" + assert captured == [ + { + "browser": "Comet", + "cookie_file": str(default_dir / "Cookies"), + "domain_name": ".xiaohongshu.com", + "os_crypt_name": "chrome", + "osx_key_service": "Comet Safe Storage", + "osx_key_user": "Comet", + } + ] + + def test_available_browsers_includes_comet(self, monkeypatch): + def fake_loader(domain_name=""): + return None + + fake_bc3 = types.SimpleNamespace(chrome=fake_loader) + monkeypatch.setitem(sys.modules, "browser_cookie3", fake_bc3) + _available_browsers.cache_clear() + + try: + assert "comet" in _available_browsers() + finally: + _available_browsers.cache_clear() + + def test_custom_browser_falls_back_to_custom_subprocess_when_no_native_loader(self, monkeypatch): + fake_bc3 = types.SimpleNamespace(chrome=lambda domain_name="": None) + monkeypatch.setitem(sys.modules, "browser_cookie3", fake_bc3) + monkeypatch.setattr( + "xhs_cli.cookies._load_custom_browser_cookies_subprocess", + lambda source, domain_name: {"a1": "custom-subprocess"}, + ) + + cookies = _extract_via_subprocess("comet") + + assert cookies == {"a1": "custom-subprocess"} + + def test_extract_browser_cookies_prefers_native_loader_when_available(self, monkeypatch): + class FakeCookie: + def __init__(self, name: str, value: str, domain: str): + self.name = name + self.value = value + self.domain = domain + + def comet(domain_name=""): + return [FakeCookie("a1", "native-a1", ".xiaohongshu.com")] + + fake_bc3 = types.SimpleNamespace(comet=comet) + monkeypatch.setitem(sys.modules, "browser_cookie3", fake_bc3) + monkeypatch.setattr( + "xhs_cli.cookies._load_custom_browser_cookies", + lambda *args, **kwargs: pytest.fail("custom fallback should not run when native loader succeeds"), + ) + + browser, cookies = extract_browser_cookies("comet") + + assert browser == "comet" + assert cookies == {"a1": "native-a1"} + class TestNoteContextCache: def test_cache_persists_token_and_source(self, tmp_config_dir): diff --git a/xhs_cli/commands/social.py b/xhs_cli/commands/social.py index 9642de6..46c2295 100644 --- a/xhs_cli/commands/social.py +++ b/xhs_cli/commands/social.py @@ -1,6 +1,7 @@ """Social commands: follow, unfollow, favorites, likes.""" -from typing import Any, Callable +from collections.abc import Callable +from typing import Any import click diff --git a/xhs_cli/cookies.py b/xhs_cli/cookies.py index 96665a4..a4a5e69 100644 --- a/xhs_cli/cookies.py +++ b/xhs_cli/cookies.py @@ -17,6 +17,22 @@ logger = logging.getLogger(__name__) +CUSTOM_BROWSER_COOKIE_CONFIGS: dict[str, dict[str, Any]] = { + "comet": { + "browser": "Comet", + "os_crypt_name": "chrome", + "osx_key_service": "Comet Safe Storage", + "osx_key_user": "Comet", + "base_dir": ("Library", "Application Support", "Comet"), + "cookie_patterns": ( + "Default/Cookies", + "Default/Network/Cookies", + "Profile */Cookies", + "Profile */Network/Cookies", + ), + } +} + # Cookie TTL: warn and attempt browser refresh after 7 days COOKIE_TTL_DAYS = 7 _COOKIE_TTL_SECONDS = COOKIE_TTL_DAYS * 86400 @@ -323,7 +339,7 @@ def _available_browsers() -> tuple[str, ...]: import browser_cookie3 as bc3 - return tuple(sorted( + standard = { name for name in dir(bc3) if not name.startswith("_") @@ -331,19 +347,165 @@ def _available_browsers() -> tuple[str, ...]: and callable(getattr(bc3, name)) and hasattr(getattr(bc3, name), "__code__") and "domain_name" in inspect.signature(getattr(bc3, name)).parameters - )) + } + return tuple(sorted(standard | set(CUSTOM_BROWSER_COOKIE_CONFIGS))) + + +def _get_custom_browser_config(source: str) -> dict[str, Any] | None: + return CUSTOM_BROWSER_COOKIE_CONFIGS.get(source.lower()) + + +def _iter_custom_browser_cookie_files(source: str) -> list[Path]: + config = _get_custom_browser_config(source) + if not config: + return [] + + base_dir = Path.home().joinpath(*config["base_dir"]) + candidates: list[Path] = [] + seen: set[Path] = set() + for pattern in config["cookie_patterns"]: + for path in sorted(base_dir.glob(pattern)): + if not path.is_file() or path in seen: + continue + seen.add(path) + candidates.append(path) + return candidates + + +def _cookies_for_domain(cookie_jar: Any, domain_name: str) -> dict[str, str]: + target_domain = domain_name.lstrip(".").lower() + + def _matches(cookie_domain: str | None) -> bool: + normalized = (cookie_domain or "").lstrip(".").lower() + return normalized == target_domain or normalized.endswith(f".{target_domain}") + + return { + cookie.name: cookie.value + for cookie in cookie_jar + if _matches(cookie.domain) + } + + +def _load_custom_browser_cookies(source: str, domain_name: str) -> dict[str, str] | None: + config = _get_custom_browser_config(source) + if not config: + return None + + try: + import browser_cookie3 as bc3 + except ImportError: + logger.debug("browser_cookie3 not installed, skipping %s extraction", source) + return None + + for cookie_file in _iter_custom_browser_cookie_files(source): + try: + loader = bc3.ChromiumBased( + browser=config["browser"], + cookie_file=str(cookie_file), + domain_name=domain_name, + os_crypt_name=config["os_crypt_name"], + osx_key_service=config["osx_key_service"], + osx_key_user=config["osx_key_user"], + ) + jar = loader.load() + except Exception as exc: + logger.debug("%s custom extraction failed for %s: %s", source, cookie_file, exc) + continue + + cookies = _cookies_for_domain(jar, domain_name) + if cookies.get("a1"): + logger.debug("Loaded XHS cookies from %s profile %s", source, cookie_file) + return cookies + + logger.debug("No usable a1 cookie found in %s custom profile scan", source) + return None + + +def _load_custom_browser_cookies_subprocess(source: str, domain_name: str) -> dict[str, str] | None: + config = _get_custom_browser_config(source) + if not config: + return None + + extract_script = ''' +import json, sys +from pathlib import Path + +try: + import browser_cookie3 as bc3 +except ImportError: + print(json.dumps({"error": "browser-cookie3 not installed"})) + sys.exit(0) + +config = json.loads(sys.argv[1]) +domain_name = sys.argv[2] +target_domain = domain_name.lstrip(".").lower() +base_dir = Path.home().joinpath(*config["base_dir"]) + +def cookies_for_domain(cookie_jar): + cookies = {} + for cookie in cookie_jar: + normalized = (cookie.domain or "").lstrip(".").lower() + if normalized == target_domain or normalized.endswith(f".{target_domain}"): + cookies[cookie.name] = cookie.value + return cookies + +for pattern in config["cookie_patterns"]: + for cookie_file in sorted(base_dir.glob(pattern)): + if not cookie_file.is_file(): + continue + try: + loader = bc3.ChromiumBased( + browser=config["browser"], + cookie_file=str(cookie_file), + domain_name=domain_name, + os_crypt_name=config["os_crypt_name"], + osx_key_service=config["osx_key_service"], + osx_key_user=config["osx_key_user"], + ) + cookies = cookies_for_domain(loader.load()) + except Exception: + continue + + if cookies.get("a1"): + print(json.dumps({"browser": config["browser"], "cookies": cookies})) + sys.exit(0) + +print(json.dumps({"error": "no_a1_cookie"})) +''' + + try: + result = subprocess.run( + [sys.executable, "-c", extract_script, json.dumps(config), domain_name], + 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 _get_browser_loader(source: str): - """Get browser cookie loader from browser_cookie3 by name.""" + """Get a native browser_cookie3 loader by name, if available.""" 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 None return loader @@ -354,27 +516,33 @@ def _extract_in_process(source: str) -> dict[str, str] | None: 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 + if loader is not None: + try: + jar = loader(domain_name=".xiaohongshu.com") + except Exception as exc: + logger.debug("%s in-process extraction failed: %s", source, exc) + else: + cookies = _cookies_for_domain(jar, ".xiaohongshu.com") + 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) + + if _get_custom_browser_config(source): + return _load_custom_browser_cookies(source, ".xiaohongshu.com") - 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.""" + try: + native_loader = _get_browser_loader(source) + except ImportError: + logger.debug("browser_cookie3 not installed, skipping subprocess extraction") + native_loader = None + extract_script = ''' import json, sys try: @@ -391,7 +559,13 @@ def _extract_via_subprocess(source: str) -> dict[str, str] | None: try: cj = loader(domain_name=".xiaohongshu.com") - cookies = {c.name: c.value for c in cj if "xiaohongshu.com" in (c.domain or "")} + target_domain = "xiaohongshu.com" + cookies = { + c.name: c.value + for c in cj + if (c.domain or "").lstrip(".").lower() == target_domain + or (c.domain or "").lstrip(".").lower().endswith(f".{target_domain}") + } if cookies.get("a1"): print(json.dumps({"browser": source, "cookies": cookies})) else: @@ -400,31 +574,33 @@ def _extract_via_subprocess(source: str) -> dict[str, str] | None: 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 + if native_loader is not None: + 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) + else: + data = json.loads(result.stdout.strip()) + if "error" in data: + logger.debug("Cookie extraction error: %s", data["error"]) + else: + return data["cookies"] + + except subprocess.TimeoutExpired: + logger.debug("Cookie extraction timed out") + except (json.JSONDecodeError, KeyError) as e: + logger.debug("Cookie extraction parse error: %s", e) + + if _get_custom_browser_config(source): + return _load_custom_browser_cookies_subprocess(source, ".xiaohongshu.com") - 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 + return None def extract_browser_cookies(source: str = "auto") -> tuple[str, dict[str, str]] | None: