scripts/_cookie_harvester.py

"""
Cookie Harvester using nodriver
================================
Dùng nodriver (Chrome CDP thuần, bypass Cloudflare + DataDome) để lấy session
cookies cho TripAdvisor, Agoda, Booking.com — sau đó tái sử dụng trong httpx
cho bulk requests mà không cần mở browser mỗi lần.

Cơ chế:
  1. Mở real Chrome (headless=False), visit từng site
  2. Chờ JS challenges giải quyết (Cloudflare 5s, DataDome auto-solve)
  3. Extract cookies: datadome (TA), cf_clearance + agoda_* (Agoda), cf_clearance (Booking)
  4. Lưu vào _COOKIE_STORE dict + file cache
  5. Auto-refresh khi cookies hết hạn (TTL: 20 phút)

Usage:
  from scripts._cookie_harvester import get_cookies, get_headers
  
  ta_headers = get_headers("tripadvisor")   # dict with Cookie header
  agoda_cookies = get_cookies("agoda")      # raw cookie string
"""

import asyncio
import base64
import json
import logging
import os
import select
import socket
import threading
import time
from pathlib import Path
from typing import Optional

logger = logging.getLogger("cookie_harvester")

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
COOKIE_TTL = 20 * 60  # 20 minutes before re-harvest
COOKIE_CACHE_FILE = Path(__file__).parent.parent / "scraped_data" / "raw_cache" / "_cookies.json"
COOKIE_CACHE_FILE.parent.mkdir(parents=True, exist_ok=True)

# Pages to harvest from each site (must trigger their bot-protection challenge)
_HARVEST_TARGETS = {
    "tripadvisor": "https://www.tripadvisor.com/Restaurants-g293924-Hanoi.html",
    "agoda":       "https://www.agoda.com/",
    "booking":     "https://www.booking.com/",
    "traveloka":   "https://www.traveloka.com/en-en/hotel",
}

# How long to wait after page load for JS challenges to resolve
_WAIT_AFTER_LOAD = {
    "tripadvisor": 10,
    "agoda":       15,   # CF Bot Management can take up to 10s
    "booking":     15,
    "traveloka":   8,
}

# ---------------------------------------------------------------------------
# Internal state
# ---------------------------------------------------------------------------
_COOKIE_STORE: dict[str, dict] = {}   # domain → {"cookie_str": "...", "harvested_at": ts}
_HARVEST_LOCK = threading.Lock()
_IS_HARVESTING = False

# User-Agent to use in harvested sessions
_UA = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
    "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36"
)

_PROXIES = [
    "http://nguyegqsY1:qZlSBJ1E@51.79.191.62:8161",   # RID ZP75087_18889 exp 08/05/2026
    "http://nguyeaYYeK:7QlkNk4y@103.187.5.219:8775",  # RID ZP75087_18888 exp 08/05/2026
]


# ---------------------------------------------------------------------------
# Cookie cache (disk)
# ---------------------------------------------------------------------------
def _load_cached_cookies() -> None:
    """Load previously harvested cookies from disk if still fresh."""
    global _COOKIE_STORE
    if not COOKIE_CACHE_FILE.exists():
        return
    try:
        data: dict = json.loads(COOKIE_CACHE_FILE.read_text(encoding="utf-8"))
        now = time.time()
        loaded = 0
        for domain, entry in data.items():
            age = now - entry.get("harvested_at", 0)
            if age < COOKIE_TTL and entry.get("cookie_str"):
                _COOKIE_STORE[domain] = entry
                loaded += 1
        if loaded:
            logger.info(f"[cookie_harvester] Loaded {loaded} cached cookie sets from disk")
    except Exception as e:
        logger.debug(f"[cookie_harvester] Could not load cache: {e}")


def _save_cookie_cache() -> None:
    """Persist current cookies to disk."""
    try:
        COOKIE_CACHE_FILE.write_text(
            json.dumps(_COOKIE_STORE, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
    except Exception as e:
        logger.debug(f"[cookie_harvester] Could not save cache: {e}")


# ---------------------------------------------------------------------------
# nodriver async harvest
# ---------------------------------------------------------------------------
# ---------------------------------------------------------------------------
# Local auth-forwarding proxy (lets Chrome use an auth proxy without dialogs)
# ---------------------------------------------------------------------------
class _LocalAuthProxy:
    """
    Minimal HTTP/CONNECT tunnel proxy that listens on localhost and forwards
    all traffic to an authenticated upstream proxy, injecting Proxy-Authorization
    automatically. Chrome can then use --proxy-server=http://127.0.0.1:PORT
    without needing to handle any auth dialog.
    """

    def __init__(self, upstream_host: str, upstream_port: int,
                 username: str, password: str, local_port: int = 18889):
        self.upstream_host = upstream_host
        self.upstream_port = upstream_port
        self._auth_header = (
            "Proxy-Authorization: Basic "
            + base64.b64encode(f"{username}:{password}".encode()).decode()
        )
        self.local_port = local_port
        self._server: Optional[socket.socket] = None
        self._thread: Optional[threading.Thread] = None
        self._running = False

    def start(self) -> int:
        self._server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self._server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._server.bind(("127.0.0.1", self.local_port))
        self._server.listen(20)
        self._server.settimeout(1.0)
        self._running = True
        self._thread = threading.Thread(target=self._accept_loop, daemon=True)
        self._thread.start()
        logger.debug(f"[local_proxy] Listening on 127.0.0.1:{self.local_port} → {self.upstream_host}:{self.upstream_port}")
        return self.local_port

    def stop(self):
        self._running = False
        if self._server:
            try:
                self._server.close()
            except Exception:
                pass

    def _accept_loop(self):
        while self._running:
            try:
                client, _ = self._server.accept()
                threading.Thread(target=self._handle_client, args=(client,), daemon=True).start()
            except socket.timeout:
                continue
            except Exception:
                break

    def _handle_client(self, client: socket.socket):
        try:
            raw = b""
            while b"\r\n\r\n" not in raw:
                chunk = client.recv(8192)
                if not chunk:
                    break
                raw += chunk
            if not raw:
                return

            head, _, body = raw.partition(b"\r\n\r\n")
            lines = head.decode(errors="replace").split("\r\n")
            first_line = lines[0]

            # Inject auth into headers
            filtered = [l for l in lines[1:] if not l.lower().startswith("proxy-authorization")]
            filtered.append(self._auth_header)
            new_head = "\r\n".join([first_line] + filtered)
            full_request = (new_head + "\r\n\r\n").encode() + body

            upstream = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            upstream.settimeout(30)
            upstream.connect((self.upstream_host, self.upstream_port))
            upstream.send(full_request)

            if first_line.upper().startswith("CONNECT"):
                # Wait for 200 Connection established
                resp = b""
                while b"\r\n\r\n" not in resp:
                    chunk = upstream.recv(4096)
                    if not chunk:
                        break
                    resp += chunk
                client.send(resp)

            self._tunnel(client, upstream)
        except Exception:
            pass
        finally:
            try:
                client.close()
            except Exception:
                pass

    @staticmethod
    def _tunnel(a: socket.socket, b: socket.socket):
        socks = [a, b]
        try:
            while True:
                r, _, _ = select.select(socks, [], socks, 10)
                if not r:
                    break
                for s in r:
                    data = s.recv(65536)
                    if not data:
                        return
                    other = b if s is a else a
                    other.sendall(data)
        except Exception:
            pass
        finally:
            for s in socks:
                try:
                    s.close()
                except Exception:
                    pass


def _parse_proxy_url(proxy_url: str):
    """Parse 'http://user:pass@host:port' → (host, port, user, pass)."""
    import urllib.parse
    p = urllib.parse.urlparse(proxy_url)
    return p.hostname, p.port, p.username or "", p.password or ""


# ---------------------------------------------------------------------------
# nodriver async harvest
# ---------------------------------------------------------------------------
async def _harvest_async(domains: list[str], proxy: Optional[str] = None) -> dict[str, str]:
    """
    Launch a real Chrome browser (via nodriver), visit each site,
    wait for JS challenges to resolve, then extract cookies.

    Returns dict: domain → cookie_string
    """
    import nodriver as uc

    # Start local auth-forwarding proxy so Chrome uses the upstream proxy
    # without needing a credentials dialog (Chrome can't do user:pass in --proxy-server)
    local_proxy: Optional[_LocalAuthProxy] = None
    effective_proxy_arg = ""
    if proxy:
        try:
            ph, pp, pu, pw = _parse_proxy_url(proxy)
            if ph and pp:
                local_proxy = _LocalAuthProxy(ph, pp, pu, pw, local_port=18889)
                local_port = local_proxy.start()
                effective_proxy_arg = f"http://127.0.0.1:{local_port}"
                logger.info(f"[cookie_harvester] Local auth proxy on :{local_port} → {ph}:{pp}")
        except Exception as e:
            logger.warning(f"[cookie_harvester] Could not start local proxy: {e}")
            local_proxy = None

    browser_args = [
        "--window-size=1280,900",
        "--lang=en-US",
        f"--user-agent={_UA}",
    ]
    if effective_proxy_arg:
        browser_args.append(f"--proxy-server={effective_proxy_arg}")
        logger.info(f"[cookie_harvester] Chrome will use proxy: {effective_proxy_arg}")
    else:
        logger.warning("[cookie_harvester] No proxy configured — direct IP may be blocked by DataDome")

    logger.info(f"[cookie_harvester] Starting Chrome for domains: {domains}")
    browser = await uc.start(headless=False, browser_args=browser_args)

    results: dict[str, str] = {}

    try:
        for domain in domains:
            url = _HARVEST_TARGETS.get(domain)
            if not url:
                continue
            wait = _WAIT_AFTER_LOAD.get(domain, 6)

            try:
                logger.info(f"[cookie_harvester] Visiting {url} ...")
                page = await browser.get(url)

                # Wait for page + CF challenge to fully resolve
                await asyncio.sleep(wait)

                # Extra wait if CF "Just a moment" still showing
                for _attempt in range(6):
                    try:
                        title = str(await page.evaluate("document.title") or "")
                        if "just a moment" in title.lower():
                            logger.info(f"[cookie_harvester] CF still active ({domain}), +3s...")
                            await asyncio.sleep(3)
                        else:
                            break
                    except Exception:
                        await asyncio.sleep(2)

                # Let nodriver settle
                await asyncio.sleep(2)

                # --- Cookie extraction ---
                # datadome is confirmed httpOnly:False in Chrome sessions — JS can read it.
                # We use JS document.cookie as primary (avoids nodriver's broken Cookie.from_json).
                # CDP fallback via raw JS JSON trick gets all same-site cookies.
                cdp_cookie_str = ""
                try:
                    result = await page.evaluate("document.cookie")
                    cdp_cookie_str = result if isinstance(result, str) else ""
                    if cdp_cookie_str:
                        logger.debug(f"[cookie_harvester] JS document.cookie: {len(cdp_cookie_str)} chars for {domain}")
                except Exception as e:
                    logger.debug(f"[cookie_harvester] JS document.cookie failed ({domain}): {e}")

                # Strategy 2: patch Cookie.from_json and use CDP (handles httpOnly if needed)
                if not cdp_cookie_str:
                    try:
                        import nodriver.cdp.network as _cdp_net
                        # Capture bound classmethod BEFORE patching
                        _orig_cfjson = _cdp_net.Cookie.from_json
                        @classmethod  # type: ignore
                        def _patched_from_json(cls, json_val):  # type: ignore
                            if isinstance(json_val, dict):
                                json_val.setdefault('sameParty', False)
                                json_val.setdefault('samesite', None)
                            return _orig_cfjson(json_val)
                        _cdp_net.Cookie.from_json = _patched_from_json
                        all_cookies = await browser.cookies.get_all()
                        _cdp_net.Cookie.from_json = _orig_cfjson  # restore
                        pairs = [f"{c.name}={c.value}" for c in all_cookies if getattr(c, 'name', None) and getattr(c, 'value', None) is not None]
                        cdp_cookie_str = "; ".join(pairs)
                        logger.debug(f"[cookie_harvester] CDP fallback got {len(all_cookies)} cookies for {domain}")
                    except Exception as e:
                        logger.debug(f"[cookie_harvester] CDP fallback failed ({domain}): {e}")

                if cdp_cookie_str:
                    results[domain] = cdp_cookie_str
                    has_cf = "cf_clearance" in cdp_cookie_str
                    has_dd = "datadome" in cdp_cookie_str
                    flags = " ".join(filter(None, [
                        "CF✓" if has_cf else "CF✗",
                        "DD✓" if has_dd else "DD✗",
                    ]))
                    logger.info(f"[cookie_harvester] ✓ {domain}: {len(cdp_cookie_str)} chars  {flags}")
                else:
                    logger.warning(f"[cookie_harvester] ✗ {domain}: no cookies captured")

            except Exception as e:
                logger.warning(f"[cookie_harvester] Failed to harvest {domain}: {e}")

    finally:
        try:
            browser.stop()
        except Exception:
            pass
        if local_proxy:
            try:
                local_proxy.stop()
            except Exception:
                pass

    return results


def _harvest_sync(domains: list[str], proxy: Optional[str] = None) -> dict[str, str]:
    """Run async harvest in a new event loop (safe to call from any thread)."""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        return loop.run_until_complete(_harvest_async(domains, proxy=proxy))
    finally:
        loop.close()


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------
def ensure_cookies(domains: Optional[list[str]] = None, force: bool = False) -> None:
    """
    Ensure cookies are available for the requested domains.
    Harvests if missing or expired. Thread-safe.

    Args:
        domains: list of "tripadvisor" | "agoda" | "booking"
                 defaults to all three
        force:   re-harvest even if still fresh
    """
    global _IS_HARVESTING
    if domains is None:
        domains = list(_HARVEST_TARGETS.keys())

    # Determine which domains actually need harvesting
    now = time.time()
    need = []
    for d in domains:
        entry = _COOKIE_STORE.get(d, {})
        age = now - entry.get("harvested_at", 0)
        if force or not entry.get("cookie_str") or age > COOKIE_TTL:
            need.append(d)

    if not need:
        return  # All fresh

    with _HARVEST_LOCK:
        # Double-check after acquiring lock
        need2 = []
        for d in need:
            entry = _COOKIE_STORE.get(d, {})
            age = now - entry.get("harvested_at", 0)
            if force or not entry.get("cookie_str") or age > COOKIE_TTL:
                need2.append(d)
        if not need2:
            return

        logger.info(f"[cookie_harvester] Harvesting cookies for: {need2}")
        # Use first proxy for cookie harvesting
        proxy = _PROXIES[0] if _PROXIES else None
        new_cookies = _harvest_sync(need2, proxy=proxy)

        ts = time.time()
        for domain, cookie_str in new_cookies.items():
            _COOKIE_STORE[domain] = {
                "cookie_str": cookie_str,
                "harvested_at": ts,
            }

        _save_cookie_cache()
        logger.info(f"[cookie_harvester] Harvest complete. Got: {list(new_cookies.keys())}")


def get_cookies(domain: str, auto_harvest: bool = True) -> Optional[str]:
    """
    Get cookie string for a domain.

    Returns cookie string like "cf_clearance=xxx; datadome=yyy" or None.
    """
    if auto_harvest:
        ensure_cookies([domain])
    entry = _COOKIE_STORE.get(domain, {})
    return entry.get("cookie_str") or None


def get_headers(domain: str, extra: Optional[dict] = None) -> dict:
    """
    Build a complete headers dict including cookies for the given domain.
    Safe to use with httpx.

    Args:
        domain: "tripadvisor" | "agoda" | "booking"
        extra:  additional headers to merge
    """
    cookies = get_cookies(domain)
    headers = {
        "User-Agent": _UA,
        "Accept-Language": "en-US,en;q=0.9",
        "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
        "Referer": _HARVEST_TARGETS.get(domain, "https://www.google.com/"),
        "Sec-Fetch-Dest": "document",
        "Sec-Fetch-Mode": "navigate",
        "Sec-Fetch-Site": "same-origin",
        "Upgrade-Insecure-Requests": "1",
    }
    if cookies:
        headers["Cookie"] = cookies
    if extra:
        headers.update(extra)
    return headers


def invalidate(domain: str) -> None:
    """Mark domain cookies as expired — triggers re-harvest on next use."""
    if domain in _COOKIE_STORE:
        _COOKIE_STORE[domain]["harvested_at"] = 0
        logger.info(f"[cookie_harvester] Invalidated cookies for {domain}")


# ---------------------------------------------------------------------------
# Init: try loading from disk on import
# ---------------------------------------------------------------------------
_load_cached_cookies()


# ---------------------------------------------------------------------------
# Standalone runner
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    print("Harvesting cookies for TripAdvisor, Agoda, Booking.com...")
    ensure_cookies(force=True)
    for domain in ("tripadvisor", "agoda", "booking"):
        c = get_cookies(domain, auto_harvest=False)
        if c:
            print(f"\n[{domain}] ({len(c)} chars):")
            # Show only key cookie names for privacy
            names = [p.split("=")[0].strip() for p in c.split(";")]
            print(f"  Cookies: {', '.join(names)}")
        else:
            print(f"\n[{domain}] FAILED — no cookies")