web_interface.py

import asyncio
import base64
from io import BytesIO
from PIL import Image

from asyncio import Semaphore
from abc import ABC, abstractmethod
from typing import List, Tuple, Optional
from playwright_stealth import stealth_async
from playwright.async_api import async_playwright, BrowserContext, Page, Playwright

from logger import Logger
from config import get_settings

logger = Logger()
config = get_settings()

def resize_image(sc: bytes):
    try:
        resize_perc = config.screenshot_resize_perc
        image = Image.open(BytesIO(sc))
        resized_image = image.resize(tuple(int(resize_perc * s) for s in image.size), Image.Resampling.LANCZOS)
        
        buffer = BytesIO()
        resized_image.save(buffer, format=image.format)
        resized_image_bytes = buffer.getvalue()
        return resized_image_bytes
    except Exception as e:
        logger.log(f"An error occurred while resizing image: {e}", level="error")
        return sc

class WebInterface(ABC):
    @abstractmethod
    async def fetch(self):
        pass

class HeadlessFetcher:
    def __init__(self, max_browsers: int = config.max_browsers, max_tabs_per_browser: int = config.max_tabs_per_browser):
        """
        Initializes the fetcher with browser and tab limits.

        Args:
            max_browsers (int): Maximum number of browser contexts to run in parallel.
            max_tabs_per_browser (int): Maximum number of tabs per browser context.
        """
        self.max_browsers = max_browsers
        self.max_tabs_per_browser = max_tabs_per_browser
        self.context_pool: List[BrowserContext] = []
        self.playwright: Optional[Playwright] = None
        
        # Semaphores control concurrency of contexts and total tabs:
        self.browser_semaphore = Semaphore(max_browsers)
        self.tab_semaphore = Semaphore(max_browsers * max_tabs_per_browser)
        
        # For simple round-robin selection of the next browser context:
        self._context_index = 0
        
        self.initialized = False

    async def initialize(self):
        """Initialize the Playwright context and create browser contexts."""
        if not self.initialized:
            self.playwright = await async_playwright().start()
            
            # Create multiple persistent contexts:
            # (You may want to use separate user_data_dir for each if needed.)
            for _ in range(self.max_browsers):
                context = await self.playwright.chromium.launch_persistent_context(
                    user_data_dir="",
                    headless=config.headless
                )
                self.context_pool.append(context)
            
            self.initialized = True

    async def shutdown(self):
        """Closes all browser contexts and cleans up resources."""
        for context in self.context_pool:
            await context.close()
        self.context_pool = []
        
        if self.playwright:
            await self.playwright.stop()
            self.playwright = None
        
        self.initialized = False

    def _get_next_context(self) -> BrowserContext:
        """
        Round-robin approach to select the next available context from the pool.
        """
        context = self.context_pool[self._context_index]
        self._context_index = (self._context_index + 1) % len(self.context_pool)
        return context

    async def _handle_request(self, request, packets):
        """Handle a request event."""
        try:
            post_data = None
            if request.method in ["POST", "PUT"]:
                try:
                    post_data = await request.post_data()
                    if isinstance(post_data, bytes):
                        try:
                            post_data = post_data.decode("utf-8")
                        except UnicodeDecodeError:
                            post_data = base64.b64encode(post_data).decode("utf-8")
                except Exception:
                    post_data = None

            packets["requests"].append({
                "url": request.url,
                "method": request.method,
                "headers": dict(request.headers) if request.headers else {},
                "post_data": post_data,
                "resource_type": request.resource_type,
            })
        except Exception as e:
            logger.log(f"Error handling request: {e}", level="error")

    async def _handle_response(self, response, packets):
        """Handle a response event."""
        try:
            body = await response.body()
        except Exception:
            body = None

        packets["responses"].append({
            "url": response.url,
            "status": response.status,
            "headers": response.headers,
            "body": body,
        })

    async def fetch(self, url: str) -> Tuple[Optional[str], dict, Optional[str]]:
        """
        Fetch a webpage using a headless browser and return the rendered HTML, packets, and screenshot.

        Args:
            url (str): The URL to fetch.

        Returns:
            Tuple[Optional[str], dict, Optional[str]]: Rendered HTML, request/response packets, and screenshot (base64).
        """
        packets = {"requests": [], "responses": []}
        
        # Acquire a browser context slot:
        await self.browser_semaphore.acquire()
        try:
            # Round-robin: pick the next context
            context = self._get_next_context()

            # Acquire a tab slot:
            await self.tab_semaphore.acquire()
            try:
                # Create a new page in the chosen context
                page = await context.new_page()
                try:
                    # Attach request/response handlers
                    page.on("request", lambda request: asyncio.create_task(self._handle_request(request, packets)))
                    page.on("response", lambda response: asyncio.create_task(self._handle_response(response, packets)))

                    # Apply stealth mode
                    await stealth_async(page)

                    # Navigate to the page
                    try:
                        await page.goto(url, wait_until="load")
                        await page.wait_for_load_state("networkidle")
                    except Exception as e:
                        logger.log(f"Error navigating to {url}: {e}", level="error")
                        return None, packets, None

                    # Attempt infinite scroll or lazy-load content
                    while True:
                        try:
                            previous_height = await page.evaluate("document.body.scrollHeight")
                            await page.evaluate("window.scrollTo(0, document.body.scrollHeight)")
                            await page.wait_for_timeout(1000)
                            current_height = await page.evaluate("document.body.scrollHeight")
                            if current_height == previous_height:
                                break
                        except Exception as e:
                            logger.log(f"Error during infinite scroll handling for {url}: {e}", level="error")
                            break

                    # Extract rendered HTML
                    try:
                        rendered_html = await page.evaluate("() => document.documentElement.outerHTML")
                    except Exception as e:
                        logger.log(f"Error extracting DOM for {url}: {e}", level="error")
                        rendered_html = None

                    # Capture screenshot
                    try:
                        screenshot_bytes = await page.screenshot(type="png", full_page=config.screenshot_full_page)
                        resized_screenshot_bytes = resize_image(screenshot_bytes)
                        base64_screenshot = base64.b64encode(resized_screenshot_bytes).decode()
                    except Exception as e:
                        logger.log(f"Error taking screenshot for {url}: {e}", level="error")
                        base64_screenshot = None

                    return rendered_html, packets, base64_screenshot

                finally:
                    # Always close the page in a finally block to avoid leaks
                    await page.close()

            except Exception as e:
                logger.log(f"Error during tab acquisition for {url}: {e}", level="error")
                return None, packets, None

            finally:
                # Release the tab semaphore
                self.tab_semaphore.release()

        except Exception as e:
            logger.log(f"Error during fetch operation for {url}: {e}", level="error")
            return None, packets, None

        finally:
            # Release the browser context semaphore
            self.browser_semaphore.release()


class UndetectableFetcher(WebInterface):
    # For dynamic sites which implement antibot mechanisms.
    pass