yt-channel-scraper/src/yt_scraper/cookies.py

from __future__ import annotations

import logging
import uuid
from datetime import datetime, timezone
from pathlib import Path

from .store import CookieRow, Store


log = logging.getLogger(__name__)

SESSION_COOKIE_NAMES = {"SID", "SAPISID", "__Secure-3PSID", "SSID", "LOGIN_INFO", "HSID", "APISID"}

_DEFAULT_DIR = Path("cookies")


def cookies_dir(override: str | Path | None = None) -> Path:
    d = Path(override) if override else _DEFAULT_DIR
    d.mkdir(parents=True, exist_ok=True)
    return d.resolve()


def parse_netscape(text: str) -> tuple[bool, dict]:
    """Parse Netscape cookie text. Returns (ok, info).

    info = {count, has_session, expires_at (ISO or None), names: set}.
    """
    lines = text.splitlines()
    expiries: list[int] = []
    names: set[str] = set()
    count = 0
    for line in lines:
        line = line.rstrip("\n")
        if not line.strip() or line.startswith("#"):
            continue
        parts = line.split("\t")
        if len(parts) < 7:
            continue
        try:
            expiry = int(parts[4])
        except ValueError:
            expiry = 0
        name = parts[5].strip()
        if not name:
            continue
        names.add(name)
        count += 1
        if expiry:
            expiries.append(expiry)
    if count == 0:
        return False, {"count": 0, "has_session": False, "expires_at": None, "names": set()}
    has_session = bool(names & SESSION_COOKIE_NAMES)
    expires_at = None
    if expiries:
        earliest = min(expiries)
        if earliest > 0:
            expires_at = datetime.fromtimestamp(earliest, tz=timezone.utc).isoformat(timespec="seconds")
    return True, {"count": count, "has_session": has_session, "expires_at": expires_at, "names": names}


def parse_netscape_file(path: str | Path) -> tuple[bool, dict]:
    p = Path(path)
    if not p.exists():
        return False, {"count": 0, "has_session": False, "expires_at": None, "names": set()}
    try:
        text = p.read_text(encoding="utf-8")
    except OSError as exc:
        log.warning("cannot read cookie file %s: %s", path, exc)
        return False, {"count": 0, "has_session": False, "expires_at": None, "names": set()}
    return parse_netscape(text)


def import_file(store: Store, path: str | Path, label: str | None = None, cookie_dir: str | Path | None = None) -> str:
    src = Path(path)
    text = src.read_text(encoding="utf-8")
    return import_text(store, text, label=label or src.stem, cookie_dir=cookie_dir)


def import_text(store: Store, text: str, label: str, cookie_dir: str | Path | None = None) -> str:
    ok, info = parse_netscape(text)
    if not ok:
        raise ValueError("Invalid Netscape cookie file: no parseable cookie lines")
    cookie_id = str(uuid.uuid4())
    cdir = cookies_dir(cookie_dir)
    out_path = cdir / f"{cookie_id}.txt"
    out_path.write_text(text, encoding="utf-8")
    store.upsert_cookie(
        cookie_id=cookie_id,
        filename=out_path.name,
        label=label,
        expires_at=info["expires_at"],
        has_session=info["has_session"],
        cookie_count=info["count"],
    )
    return cookie_id


def auto_import_dir(store: Store, dir_path: str | Path | None = None) -> int:
    """Import any loose .txt Netscape files in dir that aren't tracked yet. Returns count."""
    cdir = cookies_dir(dir_path)
    tracked = {c.filename for c in store.list_cookies()}
    n = 0
    for f in sorted(cdir.glob("*.txt")):
        if f.name in tracked:
            continue
        ok, info = parse_netscape_file(f)
        if not ok:
            continue
        cookie_id = f.stem
        # store the tracked id == filename stem; keep file in place
        store.upsert_cookie(
            cookie_id=cookie_id,
            filename=f.name,
            label=f.stem,
            expires_at=info["expires_at"],
            has_session=info["has_session"],
            cookie_count=info["count"],
        )
        n += 1
    # activate first cookie if none active
    if not store.get_active_cookie():
        cookies = store.list_cookies()
        if cookies:
            store.set_active_cookie(cookies[0].id)
    return n


def set_active(store: Store, cookie_id: str) -> None:
    store.set_active_cookie(cookie_id)


def list_vault(store: Store) -> list[CookieRow]:
    return store.list_cookies()


def delete(store: Store, cookie_id: str, cookie_dir: str | Path | None = None) -> None:
    row = store.get_cookie(cookie_id)
    store.delete_cookie(cookie_id)
    if row:
        cdir = cookies_dir(cookie_dir)
        (cdir / row.filename).unlink(missing_ok=True)


def resolve_active_path(store: Store, cookie_dir: str | Path | None = None) -> str | None:
    row = store.get_active_cookie()
    if not row:
        return None
    cdir = cookies_dir(cookie_dir)
    path = cdir / row.filename
    return str(path) if path.exists() else None


def is_expired(row: CookieRow) -> bool:
    if not row.expires_at:
        return False
    try:
        exp = datetime.fromisoformat(row.expires_at)
    except ValueError:
        return False
    return datetime.now(timezone.utc) > exp