Merge pull request #245 from teng-lin/worktree-issue-233

feat: add rookiepy browser cookie login (--browser-cookies)

Author: teng-lin

pyproject.toml

@@ -35,6 +35,7 @@ Issues = "https://github.com/teng-lin/notebooklm-py/issues"
 
 [project.optional-dependencies]
 browser = ["playwright>=1.40.0"]
+cookies = ["rookiepy>=0.1.0"]
 dev = [
     "pytest>=8.0.0",
     "pytest-asyncio>=0.23.0",

src/notebooklm/auth.py

@@ -254,6 +254,57 @@ def _is_allowed_auth_domain(domain: str) -> bool:
     return domain in ALLOWED_COOKIE_DOMAINS or _is_google_domain(domain)
 
 
+def convert_rookiepy_cookies_to_storage_state(
+    rookiepy_cookies: list[dict],
+) -> dict[str, Any]:
+    """Convert rookiepy cookie dicts to Playwright storage_state.json format.
+
+    Key mappings:
+    - ``http_only`` → ``httpOnly`` (snake_case to camelCase)
+    - ``expires=None`` → ``expires=-1`` (Playwright convention for session cookies)
+    - ``sameSite`` always ``"None"`` for cross-site Google cookies
+
+    Args:
+        rookiepy_cookies: List of cookie dicts from any ``rookiepy.*()`` call.
+            Required keys: ``domain``, ``name``, ``value``.
+
+    Returns:
+        Dict matching storage_state.json schema: ``{"cookies": [...], "origins": []}``.
+        Cookies missing required fields or from non-Google domains are silently skipped.
+    """
+    converted = []
+    for cookie in rookiepy_cookies:
+        domain = cookie.get("domain", "")
+        name = cookie.get("name", "")
+        value = cookie.get("value", "")
+
+        # Validate required fields
+        if not name or not value or not domain:
+            continue
+
+        if not _is_allowed_auth_domain(domain):
+            continue
+
+        path = cookie.get("path", "/")
+        http_only = cookie.get("http_only", False)
+        secure = cookie.get("secure", False)
+        expires = cookie.get("expires")
+
+        converted.append(
+            {
+                "name": name,
+                "value": value,
+                "domain": domain,
+                "path": path,
+                "expires": expires if expires is not None else -1,
+                "httpOnly": http_only,
+                "secure": secure,
+                "sameSite": "None",
+            }
+        )
+    return {"cookies": converted, "origins": []}
+
+
 def extract_cookies_from_storage(storage_state: dict[str, Any]) -> dict[str, str]:
     """Extract Google cookies from Playwright storage state for NotebookLM auth.
 

src/notebooklm/cli/session.py

@@ -20,9 +20,17 @@
 from typing import Any
 
 import click
+import httpx
 from rich.table import Table
 
-from ..auth import AuthTokens
+from ..auth import (
+    ALLOWED_COOKIE_DOMAINS,
+    GOOGLE_REGIONAL_CCTLDS,
+    AuthTokens,
+    convert_rookiepy_cookies_to_storage_state,
+    extract_cookies_from_storage,
+    fetch_tokens,
+)
 from ..client import NotebookLMClient
 from ..paths import (
     get_browser_profile_dir,
@@ -65,6 +73,170 @@
     "  4. Check if notebooklm.google.com is accessible in your browser"
 )
 
+# Maps user-facing browser names to rookiepy function names.
+_ROOKIEPY_BROWSER_ALIASES: dict[str, str] = {
+    "arc": "arc",
+    "brave": "brave",
+    "chrome": "chrome",
+    "chromium": "chromium",
+    "edge": "edge",
+    "firefox": "firefox",
+    "ie": "ie",
+    "librewolf": "librewolf",
+    "octo": "octo",
+    "opera": "opera",
+    "opera-gx": "opera_gx",
+    "opera_gx": "opera_gx",
+    "safari": "safari",
+    "vivaldi": "vivaldi",
+    "zen": "zen",
+}
+
+
+def _handle_rookiepy_error(e: Exception, browser_name: str) -> None:
+    """Print a user-friendly error for rookiepy exceptions."""
+    msg = str(e).lower()
+    if "lock" in msg or "database" in msg:
+        console.print(
+            f"[red]Could not read {browser_name} cookies: browser database is locked.[/red]\n"
+            "Close your browser and try again."
+        )
+    elif "permission" in msg or "access" in msg:
+        console.print(
+            f"[red]Permission denied reading {browser_name} cookies.[/red]\n"
+            "You may need to grant Terminal/Python access to your browser profile directory."
+        )
+    elif "keychain" in msg or "decrypt" in msg:
+        console.print(
+            f"[red]Could not decrypt {browser_name} cookies.[/red]\n"
+            "On macOS, allow Keychain access when prompted, or try a different browser."
+        )
+    else:
+        console.print(f"[red]Failed to read cookies from {browser_name}:[/red] {e}")
+
+
+def _login_with_browser_cookies(storage_path: Path, browser_name: str) -> None:
+    """Extract Google cookies from an installed browser via rookiepy.
+
+    Args:
+        storage_path: Where to write storage_state.json.
+        browser_name: "auto" to use rookiepy.load(), or a specific browser name.
+    """
+    try:
+        import rookiepy
+    except ImportError:
+        console.print(
+            "[red]rookiepy is not installed.[/red]\n"
+            "Install it with:\n"
+            "  pip install 'notebooklm-py[cookies]'\n"
+            "or directly:\n"
+            "  pip install rookiepy"
+        )
+        raise SystemExit(1) from None
+
+    # Build domains list including base and regional Google domains for rookiepy
+    domains = list(ALLOWED_COOKIE_DOMAINS)
+    # Add regional Google auth domains (e.g., .google.co.uk, .google.com.sg)
+    for cctld in GOOGLE_REGIONAL_CCTLDS:
+        domain = f".google.{cctld}"
+        if domain not in domains:
+            domains.append(domain)
+
+    if browser_name == "auto":
+        console.print("[yellow]Reading cookies from installed browser (auto-detect)...[/yellow]")
+        try:
+            raw_cookies = rookiepy.load(domains=domains)
+        except (OSError, RuntimeError) as e:
+            # OSError: file access issues (locked DB, permission denied)
+            # RuntimeError: decryption/keychain errors
+            _handle_rookiepy_error(e, "auto-detect")
+            raise SystemExit(1) from None
+    else:
+        canonical = _ROOKIEPY_BROWSER_ALIASES.get(browser_name.lower())
+        if canonical is None:
+            console.print(
+                f"[red]Unknown browser: '{browser_name}'[/red]\n"
+                f"Supported: {', '.join(sorted(_ROOKIEPY_BROWSER_ALIASES))}"
+            )
+            raise SystemExit(1)
+        console.print(f"[yellow]Reading cookies from {browser_name}...[/yellow]")
+        browser_fn = getattr(rookiepy, canonical, None)
+        if browser_fn is None or not callable(browser_fn):
+            console.print(
+                f"[red]rookiepy does not support '{canonical}' on this platform.[/red]\n"
+                "Check that rookiepy is properly installed: pip install rookiepy"
+            )
+            raise SystemExit(1)
+        try:
+            raw_cookies = browser_fn(domains=domains)
+        except (OSError, RuntimeError) as e:
+            # OSError: file access issues (locked DB, permission denied)
+            # RuntimeError: decryption/keychain errors
+            _handle_rookiepy_error(e, browser_name)
+            raise SystemExit(1) from None
+
+    storage_state = convert_rookiepy_cookies_to_storage_state(raw_cookies)
+    try:
+        cookies = extract_cookies_from_storage(storage_state)  # validates SID is present
+    except ValueError as e:
+        console.print(
+            "[red]No valid Google authentication cookies found.[/red]\n"
+            f"{e}\n\n"
+            "Make sure you are logged into Google in your browser."
+        )
+        raise SystemExit(1) from None
+
+    # Create parent directory (avoid mode= on Windows to prevent ACL issues)
+    try:
+        storage_path.parent.mkdir(parents=True, exist_ok=True)
+        storage_path.write_text(
+            json.dumps(storage_state, indent=2, ensure_ascii=False), encoding="utf-8"
+        )
+        if sys.platform != "win32":
+            # On Unix: ensure both directory and file have restrictive permissions
+            storage_path.parent.chmod(0o700)
+            storage_path.chmod(0o600)
+    except OSError as e:
+        logger.error("Failed to save authentication to %s: %s", storage_path, e)
+        console.print(
+            f"[red]Failed to save authentication to {storage_path}.[/red]\n" f"Details: {e}"
+        )
+        raise SystemExit(1) from None
+
+    console.print(f"\n[green]Authentication saved to:[/green] {storage_path}")
+
+    # Verify that cookies work — reuse cookies extracted above (no redundant disk read)
+    try:
+        run_async(fetch_tokens(cookies))
+        logger.info("Cookies verified successfully")
+        console.print("[green]Cookies verified successfully.[/green]")
+    except ValueError as e:
+        # Cookie validation failed - the extracted cookies are invalid
+        logger.error("Extracted cookies are invalid: %s", e)
+        console.print(
+            "[red]Warning: Extracted cookies failed validation.[/red]\n"
+            "The cookies may be expired or malformed.\n"
+            f"Error: {e}\n\n"
+            "Saved anyway, but you may need to re-run login if these are invalid."
+        )
+    except httpx.RequestError as e:
+        # Network error - can't verify but cookies might be OK
+        logger.warning("Could not verify cookies due to network error: %s", e)
+        console.print(
+            "[yellow]Warning: Could not verify cookies (network issue).[/yellow]\n"
+            "Cookies saved but may not be working.\n"
+            "Try running 'notebooklm ask' to test authentication."
+        )
+    except Exception as e:
+        # Unexpected error - log it fully
+        logger.warning("Unexpected error verifying cookies: %s: %s", type(e).__name__, e)
+        console.print(
+            f"[yellow]Warning: Unexpected error during verification: {e}[/yellow]\n"
+            "Cookies saved but please verify with 'notebooklm auth check --test'"
+        )
+
+    _sync_server_language_to_config()
+
 
 def _sync_server_language_to_config() -> None:
     """Fetch server language setting and persist to local config.
@@ -184,7 +356,19 @@ def register_session_commands(cli):
         default="chromium",
         help="Browser to use for login (default: chromium). Use 'msedge' for Microsoft Edge.",
     )
-    def login(storage, browser):
+    @click.option(
+        "--browser-cookies",
+        "browser_cookies",
+        default=None,
+        is_flag=False,
+        flag_value="auto",
+        help=(
+            "Read cookies from an installed browser instead of launching Playwright. "
+            "Optionally specify browser: chrome, firefox, brave, edge, safari, arc, ... "
+            "Requires: pip install 'notebooklm[cookies]'"
+        ),
+    )
+    def login(storage, browser, browser_cookies):
         """Log in to NotebookLM via browser.
 
         Opens a browser window for Google login. After logging in,
@@ -207,6 +391,12 @@ def login(storage, browser):
             )
             raise SystemExit(1)
 
+        # rookiepy fast-path: skip Playwright entirely
+        if browser_cookies is not None:
+            resolved_storage = Path(storage) if storage else get_storage_path()
+            _login_with_browser_cookies(resolved_storage, browser_cookies)
+            return
+
         storage_path = Path(storage) if storage else get_storage_path()
         browser_profile = get_browser_profile_dir()
         if sys.platform == "win32":