"""Symlink hardening for `notebooklm source add` file uploads. The CLI must reject symlinked paths unless the user explicitly opts in via ``--follow-symlinks``. This guards against the footgun where a workspace file like ``~/Downloads/foo.pdf`` is silently a symlink into ``/etc/`` or similar — uploading the symlink target would leak sensitive content. Tests cover both the **auto-detect** path (no ``--type``) and the **explicit** ``--type file`` path — both must enforce the same gate. The gate also rejects: - Leaf symlinks (default). - Parent-directory symlinks (``dir_link/file.pdf``). - Broken symlinks (``exists()`` returns False but ``is_symlink()`` is True — without the special-case we'd silently treat the path string as text content). When ``--follow-symlinks`` is passed, the resolved path is forwarded to ``add_file()`` so the upload opens exactly the file we validated (closes the TOCTOU window). """ from __future__ import annotations import os from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch import pytest from click.testing import CliRunner import notebooklm.auth as auth_module import notebooklm.cli.helpers as helpers_module from notebooklm import paths as paths_module from notebooklm.notebooklm_cli import cli from notebooklm.types import Source def inject_client(client, *, recorder=None): """Local copy of the ``cli/conftest`` helper. Defined inline because this test lives in ``tests/unit/`` (outside the ``tests/unit/cli`` package), so it cannot relatively import the shared helper and a cross-package absolute import would be fragile. """ def factory(auth=None, **kwargs): if recorder is not None: recorder.append((auth, kwargs)) return client return {"client_factory": factory} @pytest.fixture def runner() -> CliRunner: return CliRunner() @pytest.fixture def mock_auth(tmp_path: Path, monkeypatch): """Stub auth loading + token fetch so CLI paths run offline. Also points ``get_storage_path`` at a non-existent path inside the test's ``tmp_path`` so ``build_cookie_jar`` takes the in-memory branch (using the cookie dict returned by our mock) instead of reading any real storage file the runner may have left behind from a prior test. Without this, ``build_cookie_jar`` calls ``build_httpx_cookies_from_storage`` on the runner's default-profile storage, which validates against the real cookie set and fails this test with "Missing required cookies: SID, __Secure-1PSIDTS". """ fake_storage = tmp_path / "no_such_storage.json" monkeypatch.setattr(paths_module, "get_storage_path", lambda profile=None, **_kw: fake_storage) with ( patch.object(helpers_module, "load_auth_from_storage") as mock_load, patch.object( auth_module, "fetch_tokens_with_domains", new_callable=AsyncMock ) as mock_fetch, ): mock_load.return_value = { "SID": "test", "HSID": "test", "SSID": "test", "APISID": "test", "SAPISID": "test", } mock_fetch.return_value = ("csrf", "session") yield mock_load def _make_client() -> MagicMock: """Build a mock NotebookLMClient suitable for ``source add`` paths. Pre-seeds notebook list so ``-n nb_123`` resolves cleanly. """ client = MagicMock() client.__aenter__ = AsyncMock(return_value=client) client.__aexit__ = AsyncMock(return_value=None) for ns in ("notebooks", "sources", "artifacts", "chat", "research", "notes", "sharing"): setattr(client, ns, MagicMock()) nb_obj = MagicMock() nb_obj.id = "nb_123" nb_obj.title = "Test Notebook" client.notebooks.list = AsyncMock(return_value=[nb_obj]) client.notebooks.get = AsyncMock(return_value=nb_obj) # Default upload stubs; individual tests override as needed. client.sources.add_file = AsyncMock(return_value=Source(id="src_default", title="x")) client.sources.add_text = AsyncMock(return_value=Source(id="src_text", title="x")) return client def _invoke(runner: CliRunner, mock_client: MagicMock, argv: list[str]): """Run the CLI with ``mock_client`` injected via ``ctx.obj``.""" return runner.invoke(cli, argv, obj=inject_client(mock_client), catch_exceptions=False) class TestAutoDetectFilePath: """Auto-detect (no ``--type``) branch — the original PR surface.""" def test_plain_regular_file_uploads(self, runner, mock_auth, tmp_path: Path) -> None: """A non-symlinked file passes the gate and is uploaded. Also locks down the *exact* path forwarded to ``add_file()``: the resolved Path, not the raw argument. This is what closes the TOCTOU window (issue called out by coderabbit on the initial review pass). """ real_file = tmp_path / "doc.pdf" real_file.write_bytes(b"fake pdf content") assert not real_file.is_symlink() mock_client = _make_client() mock_client.sources.add_file = AsyncMock( return_value=Source(id="src_real", title="doc.pdf") ) result = _invoke(runner, mock_client, ["source", "add", str(real_file), "-n", "nb_123"]) assert result.exit_code == 0, result.output mock_client.sources.add_file.assert_awaited_once() called_path = mock_client.sources.add_file.await_args.args[1] # The CLI passes the resolved path so add_file opens the same # file we validated. For a non-symlink in a real tmp_path, the # resolved path equals the raw input after expanduser(). assert called_path == str(real_file.expanduser().resolve()) def test_symlink_without_flag_is_rejected(self, runner, mock_auth, tmp_path: Path) -> None: """Symlink without ``--follow-symlinks`` exits 1 with a clear error.""" target = tmp_path / "target.pdf" target.write_bytes(b"sensitive content") link = tmp_path / "link.pdf" os.symlink(target, link) assert link.is_symlink() mock_client = _make_client() mock_client.sources.add_file = AsyncMock( return_value=Source(id="should_not_be_used", title="should_not_be_used") ) result = _invoke(runner, mock_client, ["source", "add", str(link), "-n", "nb_123"]) assert result.exit_code == 1 assert "symlink" in result.output.lower() assert "--follow-symlinks" in result.output assert str(link) in result.output mock_client.sources.add_file.assert_not_awaited() def test_symlink_with_flag_is_followed_and_uploaded( self, runner, mock_auth, tmp_path: Path ) -> None: """``--follow-symlinks`` resolves the symlink and uploads the target. Verifies the *resolved* path is what's handed to ``add_file()`` (not the symlink path), closing the TOCTOU window. """ target = tmp_path / "target.pdf" target.write_bytes(b"real content") link = tmp_path / "link.pdf" os.symlink(target, link) assert link.is_symlink() mock_client = _make_client() mock_client.sources.add_file = AsyncMock( return_value=Source(id="src_resolved", title="target.pdf") ) result = _invoke( runner, mock_client, ["source", "add", str(link), "-n", "nb_123", "--follow-symlinks"], ) assert result.exit_code == 0, result.output mock_client.sources.add_file.assert_awaited_once() called_path = mock_client.sources.add_file.await_args.args[1] # The upload must hit the *target*, not the link, after opt-in. assert called_path == str(target.expanduser().resolve()) def test_broken_symlink_without_flag_is_rejected( self, runner, mock_auth, tmp_path: Path ) -> None: """Broken symlink without flag is rejected by the symlink gate. ``exists()`` follows symlinks and returns False for dangling links — without our special-case OR, the CLI would silently treat the broken-link *path string* as text content. We assert the symlink-rejection error fires. """ missing = tmp_path / "does_not_exist.pdf" link = tmp_path / "broken_link.pdf" os.symlink(missing, link) assert link.is_symlink() assert not link.exists() # broken — exists() follows the link mock_client = _make_client() mock_client.sources.add_text = AsyncMock( return_value=Source(id="should_not_be_used_text", title="broken_link.pdf") ) result = _invoke(runner, mock_client, ["source", "add", str(link), "-n", "nb_123"]) assert result.exit_code == 1 assert "symlink" in result.output.lower() assert "--follow-symlinks" in result.output # Critical: no upload of any kind. We must not silently turn the # link string into a text source. mock_client.sources.add_file.assert_not_awaited() mock_client.sources.add_text.assert_not_awaited() def test_broken_symlink_with_flag_falls_through_to_not_a_file( self, runner, mock_auth, tmp_path: Path ) -> None: """Broken symlink + ``--follow-symlinks`` hits the regular-file gate. After the symlink gate is bypassed, ``resolve()`` produces a path whose target doesn't exist, so ``is_file()`` is False and we surface "Not a regular file" instead of attempting an upload. """ missing = tmp_path / "does_not_exist.pdf" link = tmp_path / "broken_link.pdf" os.symlink(missing, link) assert link.is_symlink() assert not link.exists() mock_client = _make_client() mock_client.sources.add_file = AsyncMock( return_value=Source(id="should_not_be_used", title="should_not_be_used") ) result = _invoke( runner, mock_client, ["source", "add", str(link), "-n", "nb_123", "--follow-symlinks"], ) assert result.exit_code == 1 assert "not a regular file" in result.output.lower() mock_client.sources.add_file.assert_not_awaited() def test_parent_directory_symlink_is_rejected(self, runner, mock_auth, tmp_path: Path) -> None: """A symlinked *parent* directory triggers the gate. ``raw.is_symlink()`` only inspects the leaf path, so ``dir_link/file.pdf`` would slip through a leaf-only check. The helper walks parents — assert here that the gate fires. """ real_dir = tmp_path / "real_dir" real_dir.mkdir() real_file = real_dir / "doc.pdf" real_file.write_bytes(b"content") dir_link = tmp_path / "dir_link" os.symlink(real_dir, dir_link) path_through_link = dir_link / "doc.pdf" assert path_through_link.exists() # the leaf itself is a real file assert not path_through_link.is_symlink() # leaf alone looks clean assert dir_link.is_symlink() # but the parent is a link mock_client = _make_client() result = _invoke( runner, mock_client, ["source", "add", str(path_through_link), "-n", "nb_123"] ) assert result.exit_code == 1 assert "symlink" in result.output.lower() mock_client.sources.add_file.assert_not_awaited() class TestExplicitTypeFile: """``--type file`` path — the security gap closed in the review pass. Without these tests, ``source add link.pdf --type file`` would silently bypass the symlink gate and leak the link target. Locked down by parametrized variants that mirror the auto-detect tests. """ def test_explicit_type_file_with_plain_file_uploads( self, runner, mock_auth, tmp_path: Path ) -> None: """Sanity check: ``--type file`` on a normal file still works.""" real_file = tmp_path / "doc.pdf" real_file.write_bytes(b"fake pdf content") mock_client = _make_client() mock_client.sources.add_file = AsyncMock( return_value=Source(id="src_real", title="doc.pdf") ) result = _invoke( runner, mock_client, ["source", "add", str(real_file), "--type", "file", "-n", "nb_123"], ) assert result.exit_code == 0, result.output mock_client.sources.add_file.assert_awaited_once() called_path = mock_client.sources.add_file.await_args.args[1] assert called_path == str(real_file.expanduser().resolve()) def test_explicit_type_file_with_symlink_no_flag_is_rejected( self, runner, mock_auth, tmp_path: Path ) -> None: """Critical: explicit ``--type file`` MUST honor the symlink gate. This is the path most likely used by automation scripts (which often hardcode ``--type``), so it's the higher-risk surface. Pre-fix, this case bypassed the guard completely. """ target = tmp_path / "target.pdf" target.write_bytes(b"sensitive content") link = tmp_path / "link.pdf" os.symlink(target, link) mock_client = _make_client() mock_client.sources.add_file = AsyncMock( return_value=Source(id="should_not_be_used", title="should_not_be_used") ) result = _invoke( runner, mock_client, ["source", "add", str(link), "--type", "file", "-n", "nb_123"], ) assert result.exit_code == 1 assert "symlink" in result.output.lower() assert "--follow-symlinks" in result.output mock_client.sources.add_file.assert_not_awaited() def test_explicit_type_file_with_symlink_and_flag_is_followed( self, runner, mock_auth, tmp_path: Path ) -> None: """``--type file --follow-symlinks`` resolves and uploads.""" target = tmp_path / "target.pdf" target.write_bytes(b"real content") link = tmp_path / "link.pdf" os.symlink(target, link) mock_client = _make_client() mock_client.sources.add_file = AsyncMock( return_value=Source(id="src_resolved", title="target.pdf") ) result = _invoke( runner, mock_client, [ "source", "add", str(link), "--type", "file", "-n", "nb_123", "--follow-symlinks", ], ) assert result.exit_code == 0, result.output mock_client.sources.add_file.assert_awaited_once() called_path = mock_client.sources.add_file.await_args.args[1] assert called_path == str(target.expanduser().resolve())