yt-channel-scraper/tests/test_cookies.py

from __future__ import annotations

from pathlib import Path
import pytest

from yt_scraper.cookies import (
    BrowserCookieLockedError,
    auto_import_dir,
    delete,
    import_from_browser,
    is_expired,
    parse_netscape,
    resolve_active_path,
)
from yt_scraper.store import CookieRow, Store


SAMPLE_NETSCAPE = """# Netscape HTTP Cookie File
.youtube.com\tTRUE\t/\tTRUE\t2000000000\tSID\tsample_sid_token
.youtube.com\tTRUE\t/\tTRUE\t2000000000\tLOGIN_INFO\tsample_login_info
"""


def test_parse_netscape():
    ok, info = parse_netscape(SAMPLE_NETSCAPE)
    assert ok is True
    assert info["count"] == 2
    assert info["has_session"] is True
    assert info["expires_at"] is not None


def test_parse_netscape_httponly_lines_are_data():
    # SID/HSID are HttpOnly; Netscape exports prefix those lines and a
    # comment-skipping parser would silently strip the login out.
    text = (
        "# Netscape HTTP Cookie File\n"
        "#HttpOnly_.youtube.com\tTRUE\t/\tTRUE\t2000000000\tSID\ttok\n"
        "#HttpOnly_.youtube.com\tTRUE\t/\tTRUE\t2000000000\tHSID\ttok\n"
        "#HttpOnly_.youtube.com\tTRUE\t/\tTRUE\t2000000000\tSSID\ttok\n"
        "# a real comment\n"
    )
    ok, info = parse_netscape(text)
    assert ok is True
    assert info["count"] == 3
    assert info["has_session"] is True


def test_partial_session_export_is_not_a_session():
    # A lone __Secure-3PSID (partial extension export) is anonymous to
    # YouTube; it must not be reported as a usable session.
    text = (
        "# Netscape HTTP Cookie File\n"
        ".youtube.com\tTRUE\t/\tTRUE\t2000000000\t__Secure-3PSID\ttok\n"
        ".youtube.com\tTRUE\t/\tTRUE\t2000000000\t__Secure-3PAPISID\ttok\n"
    )
    ok, info = parse_netscape(text)
    assert ok is True
    assert info["has_session"] is False


def _browser_cookie(name, value, *, domain=".youtube.com", httponly=False):
    import http.cookiejar
    return http.cookiejar.Cookie(
        version=0, name=name, value=value, port=None, port_specified=False,
        domain=domain, domain_specified=True, domain_initial_dot=domain.startswith("."),
        path="/", path_specified=True, secure=True, expires=2000000000,
        discard=False, comment=None, comment_url=None,
        rest={"HttpOnly": None} if httponly else {},
    )


def test_import_from_browser_writes_vault_file(tmp_path: Path, monkeypatch):
    store = Store(tmp_path / "state.db")
    cookie_dir = tmp_path / "cookies"

    def fake_extract(browser, profile=None, logger=None, **kw):
        return iter([
            _browser_cookie("SID", "sid_tok", httponly=True),
            _browser_cookie("HSID", "hsid_tok", httponly=True),
            _browser_cookie("SSID", "ssid_tok", httponly=True),
            _browser_cookie("__Secure-3PSID", "psid_tok"),
            _browser_cookie("PREF", "pref_tok", domain=".google.com"),  # not youtube -> dropped
        ])

    import yt_dlp.cookies as ydl_cookies
    monkeypatch.setattr(ydl_cookies, "extract_cookies_from_browser", fake_extract)

    cid = import_from_browser(store, browser="brave", cookie_dir=cookie_dir)
    row = store.get_cookie(cid)
    assert row is not None
    assert row.cookie_count == 4
    assert row.has_session is True

    # The written file must round-trip as a valid Netscape file whose
    # HttpOnly session cookies survive the vault's own parser.
    from yt_scraper.cookies import parse_netscape_file
    path = cookie_dir / row.filename
    ok, info = parse_netscape_file(path)
    assert ok is True
    assert info["count"] == 4
    assert info["has_session"] is True


def test_import_from_browser_maps_locked_db(tmp_path: Path, monkeypatch):
    store = Store(tmp_path / "state.db")

    def fake_extract(browser, profile=None, logger=None, **kw):
        raise RuntimeError("Could not copy Chrome cookie database. See  https://github.com/yt-dlp/yt-dlp/issues/7271  for more info")

    import yt_dlp.cookies as ydl_cookies
    monkeypatch.setattr(ydl_cookies, "extract_cookies_from_browser", fake_extract)

    with pytest.raises(BrowserCookieLockedError) as exc:
        import_from_browser(store, browser="brave", cookie_dir=tmp_path / "cookies")
    assert "close brave" in str(exc.value).lower()


def test_auto_import_and_prune_dead_cookies(tmp_path: Path):
    db_path = tmp_path / "state.db"
    store = Store(db_path)
    cookie_dir = tmp_path / "cookies"
    cookie_dir.mkdir()

    # Place two cookie files
    f1 = cookie_dir / "c1.txt"
    f1.write_text(SAMPLE_NETSCAPE, encoding="utf-8")
    f2 = cookie_dir / "c2.txt"
    f2.write_text(SAMPLE_NETSCAPE, encoding="utf-8")

    # auto_import imports both and activates the first
    n = auto_import_dir(store, dir_path=cookie_dir)
    assert n == 2
    assert len(store.list_cookies()) == 2
    active = store.get_active_cookie()
    assert active is not None
    assert active.filename in ("c1.txt", "c2.txt")

    active_path = resolve_active_path(store, cookie_dir=cookie_dir)
    assert active_path is not None
    assert Path(active_path).exists()

    # Now simulate user deleting the active cookie file from disk
    Path(active_path).unlink()

    # resolve_active_path should fall back to the surviving cookie file
    fallback_path = resolve_active_path(store, cookie_dir=cookie_dir)
    assert fallback_path is not None
    assert Path(fallback_path).exists()

    # auto_import_dir should clean up the deleted cookie row from DB
    auto_import_dir(store, dir_path=cookie_dir)
    assert len(store.list_cookies()) == 1