iptv/M3U8/scrapers/utils/network.py

import asyncio
import logging
import re
from collections.abc import Callable
from typing import Any

import httpx
from playwright.async_api import Browser, BrowserContext, Playwright, Request

from .config import UA

CLIENT = httpx.AsyncClient(
    timeout=5,
    follow_redirects=True,
    headers={"User-Agent": UA},
)


async def check_status(client: httpx.AsyncClient, url: str) -> bool:
    try:
        r = await client.get(url)
        r.raise_for_status()
    except Exception:
        return False

    return r.status_code == 200


async def get_base(client: httpx.AsyncClient, mirrors: list[str]) -> str | None:
    tasks = [check_status(client, link) for link in mirrors]
    results = await asyncio.gather(*tasks)

    try:
        return [url for url, ok in zip(mirrors, results) if ok][0]
    except IndexError:
        return


async def safe_process_event(
    fn: Callable,
    url_num: int,
    timeout: int | float = 15,
    log: logging.Logger | None = None,
) -> Any | None:

    if not log:
        log = logging.getLogger(__name__)

    task = asyncio.create_task(fn())

    try:
        return await asyncio.wait_for(task, timeout=timeout)
    except asyncio.TimeoutError:
        log.warning(f"URL {url_num}) Timed out after {timeout}s, skipping event")

        task.cancel()

        try:
            await task
        except asyncio.CancelledError:
            pass
        except Exception as e:
            log.debug(f"URL {url_num}) Ignore exception after timeout: {e}")


def capture_req(
    req: Request,
    captured: list[str],
    got_one: asyncio.Event,
) -> None:
    valid_m3u8 = re.compile(r"^(?!.*(amazonaws|knitcdn)).*\.m3u8")

    if valid_m3u8.search(req.url):
        captured.append(req.url)
        got_one.set()


async def new_browser(
    playwright: Playwright,
    browser: str = "firefox",
    ignore_https_errors: bool = False,
) -> tuple[Browser, BrowserContext]:

    if browser == "brave":
        brwsr = await playwright.chromium.connect_over_cdp("http://localhost:9222")
        context = brwsr.contexts[0]
    else:
        brwsr = await playwright.firefox.launch(headless=True)

        context = await brwsr.new_context(
            user_agent=UA,
            ignore_https_errors=ignore_https_errors,
            viewport={"width": 1366, "height": 768},
            device_scale_factor=1,
            locale="en-US",
            timezone_id="America/New_York",
            color_scheme="dark",
            permissions=["geolocation"],
            extra_http_headers={
                "Accept-Language": "en-US,en;q=0.9",
                "Upgrade-Insecure-Requests": "1",
            },
        )

        await context.add_init_script(
            """
        Object.defineProperty(navigator, 'webdriver', {get: () => undefined});

        Object.defineProperty(navigator, 'languages', {
          get: () => ['en-US', 'en']
        });

        Object.defineProperty(navigator, 'plugins', {
          get: () => [1, 2, 3, 4]
        });

        const elementDescriptor = Object.getOwnPropertyDescriptor(HTMLElement.prototype, 'offsetHeight');
        Object.defineProperty(HTMLDivElement.prototype, 'offsetHeight', {
          ...elementDescriptor,
          get: function() {
            if (this.id === 'modernizr') { return 24; }
            return elementDescriptor.get.apply(this);
          }
        });

        Object.defineProperty(window.screen, 'width', { get: ()     => 1366 });
        Object.defineProperty(window.screen, 'height', { get: ()    => 768 });

        const getParameter = WebGLRenderingContext.prototype.   getParameter;
        WebGLRenderingContext.prototype.getParameter = function (param) {
          if (param === 37445) return "Intel Inc."; //  UNMASKED_VENDOR_WEBGL
          if (param === 37446) return "Intel Iris OpenGL    Engine"; // UNMASKED_RENDERER_WEBGL
          return getParameter.apply(this, [param]);
        };

        const observer = new MutationObserver(mutations => {
            mutations.forEach(mutation => {
                mutation.addedNodes.forEach(node => {
                    if (node.tagName === 'IFRAME' && node.hasAttribute('sandbox')) {
                        node.removeAttribute('sandbox');
                    }
                });
            });
        });

        observer.observe(document.documentElement, { childList: true, subtree: true });
        """
        )

    return brwsr, context