browser_cookie_extractor.py

from collections.abc import Callable
from enum import Enum
from http.cookiejar import CookieJar
from typing import Any, cast

import browser_cookie3


class SupportedBrowser(Enum):
    """Supported browsers for cookie extraction."""

    CHROME = "chrome"
    FIREFOX = "firefox"
    OPERA = "opera"
    OPERA_GX = "opera gx"
    EDGE = "edge"
    CHROMIUM = "chromium"
    BRAVE = "brave"
    VIVALDI = "vivaldi"
    SAFARI = "safari"


def get_cookie_loader(
    browser: SupportedBrowser | None,
) -> Callable[[], CookieJar]:
    """Get a cookie loader for a specific browser, or the default browser."""

    browser_function_map = {
        SupportedBrowser.CHROME: browser_cookie3.chrome,
        SupportedBrowser.FIREFOX: browser_cookie3.firefox,
        SupportedBrowser.OPERA: browser_cookie3.opera,
        SupportedBrowser.OPERA_GX: browser_cookie3.opera_gx,
        SupportedBrowser.EDGE: browser_cookie3.edge,
        SupportedBrowser.CHROMIUM: browser_cookie3.chromium,
        SupportedBrowser.BRAVE: browser_cookie3.brave,
        SupportedBrowser.VIVALDI: browser_cookie3.vivaldi,
        SupportedBrowser.SAFARI: browser_cookie3.safari,
    }

    def wrapper() -> CookieJar:
        if browser:
            return browser_function_map.get(browser, browser_cookie3.load)()
        return browser_cookie3.load()

    return wrapper


def load_cookies(
    browser: SupportedBrowser | None = None,
    domain: str = "",
) -> CookieJar:
    """Load cookies from a browser, or the default browser. Optionally filter by domain."""

    cookie_loader = get_cookie_loader(browser=browser)
    all_cookies = cast(CookieJar, cookie_loader())
    if not domain:
        return all_cookies

    filtered_jar = CookieJar()
    for cookie in all_cookies:
        if cookie.domain.endswith(f"{domain}"):
            filtered_jar.set_cookie(cookie)
    return filtered_jar

# utils

def convert_cookies_for_header(cookies: CookieJar) -> str:
    """Convert cookies to a header string, e.g. 'name=value; name2=value2'"""

    return "; ".join([f"{c.name}={c.value}" for c in cookies])


def convert_cookies_to_dict(cookies: CookieJar) -> list[dict[str, Any]]:
    """Convert cookies to a list of dictionaries, e.g. [{'name': 'name', 'value': 'value'}, ...]"""

    cookies_dict: list[dict[str, Any]] = []

    for cookie in cookies:
        cookie_attrs = vars(cookie)
        target_playwright_cookie = {
            "name": cookie.name,
            "value": cookie.value,
            "domain": cookie.domain,
            "path": cookie.path,
            "httpOnly": cookie_attrs.get("_rest", {}).get("HttpOnly", False),
            "secure": bool(cookie.secure),
        }
        if cookie.expires is not None:
            target_playwright_cookie["expires"] = cookie.expires

        cookies_dict.append(target_playwright_cookie)

    return cookies_dict