ai-browser-profile 1.0.5 → 1.0.7

package/ai_browser_profile/cookies.py

@@ -0,0 +1,330 @@
+"""Read cookies from a Chromium browser profile and inject into another.
+
+Sibling to the ingestors in ai_browser_profile.ingestors. Lives in the
+package but is NOT wired into extract_memories() — cookies are auth
+secrets and must never land in memories.db.
+
+Public API:
+    read_cookies(profile, domains=None)   -> list[Cookie]
+    inject_via_cdp(cookies, cdp_url, ...) -> int
+
+CLI:
+    python -m ai_browser_profile.cookies copy \\
+        --from chrome:Default \\
+        --to cdp://127.0.0.1:9555 \\
+        --domains github.com,linear.app
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import logging
+import shutil
+import sqlite3
+import subprocess
+import sys
+import urllib.request
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable, Optional
+
+from ai_browser_profile.ingestors.browser_detect import (
+    BrowserProfile,
+    copy_db,
+    detect_browsers,
+)
+
+log = logging.getLogger(__name__)
+
+KEYCHAIN_SERVICE = {
+    "chrome": "Chrome Safe Storage",
+    "arc": "Arc Safe Storage",
+    "brave": "Brave Safe Storage",
+    "edge": "Microsoft Edge Safe Storage",
+    "chromium": "Chromium Safe Storage",
+}
+
+PBKDF2_SALT = b"saltysalt"
+PBKDF2_ITERATIONS = 1003
+AES_KEY_LENGTH = 16
+AES_IV = b" " * 16
+
+SAMESITE_MAP = {-1: "Unspecified", 0: "None", 1: "Lax", 2: "Strict"}
+
+
+@dataclass
+class Cookie:
+    name: str
+    value: str
+    domain: str
+    path: str
+    expires: float
+    secure: bool
+    http_only: bool
+    same_site: str
+
+
+def _keychain_password(browser: str) -> bytes:
+    service = KEYCHAIN_SERVICE.get(browser)
+    if not service:
+        raise ValueError(f"No keychain service mapped for browser {browser!r}")
+    res = subprocess.run(
+        ["security", "find-generic-password", "-w", "-s", service],
+        capture_output=True, text=True, check=False,
+    )
+    if res.returncode != 0:
+        raise RuntimeError(
+            f"Could not read {service!r} from Keychain: {res.stderr.strip() or 'access denied'}"
+        )
+    return res.stdout.strip().encode()
+
+
+def _derive_key(password: bytes) -> bytes:
+    return hashlib.pbkdf2_hmac(
+        "sha1", password, PBKDF2_SALT, PBKDF2_ITERATIONS, AES_KEY_LENGTH
+    )
+
+
+def _decrypt(encrypted: bytes, key: bytes, host_key: str) -> Optional[str]:
+    """Decrypt a Chromium cookie value. Returns None on failure."""
+    from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+
+    if not encrypted:
+        return None
+    prefix = encrypted[:3]
+    if prefix in (b"v10", b"v11"):
+        payload = encrypted[3:]
+    else:
+        payload = encrypted
+    if len(payload) % 16 != 0:
+        return None
+    cipher = Cipher(algorithms.AES(key), modes.CBC(AES_IV))
+    dec = cipher.decryptor()
+    plain = dec.update(payload) + dec.finalize()
+    if not plain:
+        return None
+    pad = plain[-1]
+    if 1 <= pad <= 16 and plain.endswith(bytes([pad]) * pad):
+        plain = plain[:-pad]
+    # Chrome 80+ prepends SHA256(host_key) (32 bytes) to bind cookie to its host.
+    expected = hashlib.sha256(host_key.encode()).digest()
+    if plain.startswith(expected):
+        plain = plain[32:]
+    try:
+        return plain.decode("utf-8")
+    except UnicodeDecodeError:
+        return plain.decode("utf-8", errors="replace")
+
+
+def read_cookies(
+    profile: BrowserProfile,
+    domains: Optional[Iterable[str]] = None,
+) -> list[Cookie]:
+    """Read and decrypt cookies from a Chromium browser profile.
+
+    Args:
+        profile: A Chromium profile from detect_browsers().
+        domains: Optional iterable of substrings; a cookie is kept if its
+                 host_key contains any of them. None means all cookies.
+    """
+    if profile.browser in ("safari", "firefox"):
+        raise NotImplementedError(f"Cookie read not supported for {profile.browser}")
+
+    cookies_path = profile.path / "Cookies"
+    if not cookies_path.exists():
+        raise FileNotFoundError(f"No Cookies file at {cookies_path}")
+
+    tmp = copy_db(cookies_path)
+    if tmp is None:
+        raise RuntimeError(
+            f"Could not copy {cookies_path}. Grant Full Disk Access to your terminal and retry."
+        )
+
+    domain_filters = list(domains) if domains else None
+    key = _derive_key(_keychain_password(profile.browser))
+    cookies: list[Cookie] = []
+    skipped = 0
+    def _txt(b) -> str:
+        if b is None:
+            return ""
+        if isinstance(b, bytes):
+            return b.decode("utf-8", errors="replace")
+        return str(b)
+
+    try:
+        conn = sqlite3.connect(f"file:{tmp}?mode=ro", uri=True)
+        # Arc and some Chrome forks declare encrypted_value as TEXT, not BLOB,
+        # which makes sqlite3 try to UTF-8-decode the AES ciphertext and crash
+        # mid-iteration. Force everything to bytes and decode TEXT columns
+        # ourselves.
+        conn.text_factory = bytes
+        conn.row_factory = sqlite3.Row
+        rows = conn.execute(
+            "SELECT host_key, name, value, encrypted_value, path, expires_utc, "
+            "is_secure, is_httponly, samesite FROM cookies"
+        )
+        for row in rows:
+            host = _txt(row["host_key"])
+            if domain_filters and not any(d in host for d in domain_filters):
+                continue
+            value = _txt(row["value"])
+            if not value and row["encrypted_value"]:
+                value = _decrypt(row["encrypted_value"], key, host) or ""
+                if not value:
+                    skipped += 1
+                    continue
+            expires = 0.0
+            if row["expires_utc"]:
+                # Chromium epoch is 1601-01-01 in microseconds.
+                expires = (row["expires_utc"] / 1_000_000) - 11644473600
+            cookies.append(Cookie(
+                name=_txt(row["name"]),
+                value=value,
+                domain=host,
+                path=_txt(row["path"]) or "/",
+                expires=expires,
+                secure=bool(row["is_secure"]),
+                http_only=bool(row["is_httponly"]),
+                same_site=SAMESITE_MAP.get(row["samesite"], "Unspecified"),
+            ))
+        conn.close()
+    finally:
+        shutil.rmtree(tmp.parent, ignore_errors=True)
+
+    log.info(
+        "Read %d cookies from %s/%s (skipped %d undecryptable)",
+        len(cookies), profile.browser, profile.name, skipped,
+    )
+    return cookies
+
+
+def _ws_from_cdp_url(cdp_url: str) -> str:
+    if cdp_url.startswith("ws://") or cdp_url.startswith("wss://"):
+        return cdp_url
+    if cdp_url.startswith("cdp://"):
+        cdp_url = "http://" + cdp_url[len("cdp://"):]
+    base = cdp_url.rstrip("/")
+    with urllib.request.urlopen(f"{base}/json/version", timeout=5) as r:
+        return json.loads(r.read())["webSocketDebuggerUrl"]
+
+
+def inject_via_cdp(
+    cookies: Iterable[Cookie],
+    cdp_url: str = "http://127.0.0.1:9222",
+) -> int:
+    """Inject cookies into a running Chrome via CDP Storage.setCookies.
+
+    Args:
+        cookies: iterable of Cookie objects.
+        cdp_url: base http(s) URL of the Chrome DevTools endpoint, or a
+                 cdp://host:port shorthand, or a raw ws:// URL.
+
+    Returns: number of cookies submitted (the browser accepts them as a batch).
+    """
+    from websocket import create_connection
+
+    ws_url = _ws_from_cdp_url(cdp_url)
+    # Chrome 111+ enforces CDP origin checking and rejects any Origin header
+    # unless the target was launched with --remote-allow-origins. Suppressing
+    # the header bypasses the check; localhost CDP is already privileged.
+    ws = create_connection(ws_url, timeout=10, suppress_origin=True)
+    try:
+        batch = []
+        for c in cookies:
+            param = {
+                "name": c.name,
+                "value": c.value,
+                "domain": c.domain,
+                "path": c.path or "/",
+                "secure": c.secure,
+                "httpOnly": c.http_only,
+            }
+            if c.same_site in ("Strict", "Lax", "None"):
+                param["sameSite"] = c.same_site
+            if c.expires > 0:
+                param["expires"] = c.expires
+            batch.append(param)
+        if not batch:
+            return 0
+        ws.send(json.dumps({
+            "id": 1,
+            "method": "Storage.setCookies",
+            "params": {"cookies": batch},
+        }))
+        resp = json.loads(ws.recv())
+        if "error" in resp:
+            log.warning("Storage.setCookies failed: %s", resp["error"])
+            return 0
+    finally:
+        ws.close()
+    log.info("Injected %d cookies via CDP", len(batch))
+    return len(batch)
+
+
+# --- helpers used by CLI and external callers ---
+
+def find_profile(spec: str) -> BrowserProfile:
+    """Resolve a 'browser:profile-name' spec (e.g. 'chrome:Default') to a BrowserProfile."""
+    if ":" in spec:
+        browser, name = spec.split(":", 1)
+    else:
+        browser, name = spec, "Default"
+    matches = [p for p in detect_browsers({browser}) if p.name == name]
+    if not matches:
+        available = [(p.browser, p.name) for p in detect_browsers({browser})]
+        raise SystemExit(
+            f"No profile {spec!r}. Available {browser} profiles: {available}"
+        )
+    return matches[0]
+
+
+def _cli(argv: Optional[list[str]] = None) -> int:
+    parser = argparse.ArgumentParser(prog="python -m ai_browser_profile.cookies")
+    sub = parser.add_subparsers(dest="cmd", required=True)
+
+    cp = sub.add_parser("copy", help="copy cookies from a local profile into a running browser via CDP")
+    cp.add_argument("--from", dest="src", required=True,
+                    help="source profile, e.g. chrome:Default or arc:'Profile 1'")
+    cp.add_argument("--to", dest="dst", required=True,
+                    help="target CDP endpoint, e.g. cdp://127.0.0.1:9555 or http://127.0.0.1:9555")
+    cp.add_argument("--domains", default=None,
+                    help="comma-separated list of host_key substrings to include")
+    cp.add_argument("-v", "--verbose", action="store_true")
+
+    ls = sub.add_parser("list", help="list cookies in a local profile (counts only — no values printed)")
+    ls.add_argument("--from", dest="src", required=True)
+    ls.add_argument("--domains", default=None)
+
+    args = parser.parse_args(argv)
+
+    logging.basicConfig(
+        level=logging.DEBUG if getattr(args, "verbose", False) else logging.INFO,
+        format="%(levelname)s %(message)s",
+    )
+
+    profile = find_profile(args.src)
+    domain_filters = [d.strip() for d in args.domains.split(",")] if args.domains else None
+
+    cookies = read_cookies(profile, domains=domain_filters)
+
+    if args.cmd == "list":
+        by_host: dict[str, int] = {}
+        for c in cookies:
+            by_host[c.domain] = by_host.get(c.domain, 0) + 1
+        for host, n in sorted(by_host.items(), key=lambda kv: -kv[1]):
+            print(f"  {n:4}  {host}")
+        print(f"Total: {len(cookies)} cookies across {len(by_host)} hosts")
+        return 0
+
+    if args.cmd == "copy":
+        n = inject_via_cdp(cookies, args.dst)
+        print(f"Injected {n}/{len(cookies)} cookies into {args.dst}")
+        return 0 if n > 0 else 2
+
+    return 1
+
+
+if __name__ == "__main__":
+    sys.exit(_cli())

package/ai_browser_profile/extract.py

@@ -40,6 +40,7 @@ def extract_memories(memories_db_path: str = "memories.db",
     """
     total_start = time.monotonic()
     mem = MemoryDB(memories_db_path, defer_embeddings=True)
+    from ai_browser_profile.ingestors.browser_detect import permission_denied_paths
     profiles = detect_browsers(allowed=browsers)
     log.info(f"Extracting memories from {len(profiles)} profiles...")
 
@@ -80,6 +81,17 @@ def extract_memories(memories_db_path: str = "memories.db",
     run_cleanup(db_path=memories_db_path)
     mem = MemoryDB(memories_db_path, defer_embeddings=True)
     interim_profile = mem.profile_text()
+
+    # Emit structured browser summary so callers can show transparency
+    detected_browsers = sorted(set(p.browser for p in profiles))
+    denied_browsers = sorted(set(
+        p.browser for p in profiles
+        if any(str(p.path) in str(denied) for denied in permission_denied_paths)
+    ))
+    print(f"BROWSERS_SCANNED: {','.join(detected_browsers)}", flush=True)
+    if denied_browsers:
+        print(f"BROWSERS_PERMISSION_DENIED: {','.join(denied_browsers)}", flush=True)
+
     log.info(f"Interim profile ready (WhatsApp + embeddings still running):\n{interim_profile}")
 
     # 7. WhatsApp — contacts from IndexedDB (slow, runs last)