crawler-intra-mart/extractor/asset_downloader.py

import asyncio
import hashlib
from pathlib import Path
from urllib.parse import urljoin, urlparse

import aiofiles
import httpx
from bs4 import BeautifulSoup

from utils.logger import logger
import os


def _ensure_soup(html_or_soup):
    if hasattr(html_or_soup, "select_one") and hasattr(
        html_or_soup, "find_all"
    ):
        return html_or_soup
    return BeautifulSoup(html_or_soup, "lxml")

class AssetDownloader:
    """
    Download assets from documentation pages.

    Supported:
    - images
    - css
    - js
    - fonts
    - downloadable files

    Features:
    - deduplicate downloads
    - async download
    - preserve structure
    """

    def __init__(
        self,
        output_dir="output/assets",
        timeout=30,
        max_concurrency=10,
        http_client=None
    ):
        self.output_dir = Path(output_dir)

        self.timeout = timeout

        # downloaded maps url -> local Path so repeat lookups
        # short-circuit without touching the filesystem
        self.downloaded = {}

        # Shared httpx client (keep-alive + optional http2). If
        # the caller didn't pass one we create our own — but
        # a shared client across pages is what actually buys
        # the speedup.
        self._owns_client = http_client is None
        if http_client is None:
            try:
                import h2  # noqa: F401
                _http2 = True
            except ImportError:
                _http2 = False
            http_client = httpx.AsyncClient(
                timeout=self.timeout,
                follow_redirects=True,
                http2=_http2,
                limits=httpx.Limits(
                    max_connections=max_concurrency * 2,
                    max_keepalive_connections=max_concurrency * 2,
                ),
            )
        self.client = http_client

        self.semaphore = asyncio.Semaphore(max_concurrency)

        self.allowed_extensions = {
            ".png",
            ".jpg",
            ".jpeg",
            ".gif",
            ".svg",
            ".webp",
            ".css",
            ".js",
            ".woff",
            ".woff2",
            ".ttf",
            ".eot",
            ".pdf",
            ".zip",
            ".ico",
            ".bmp",
            ".avif",
            ".map",
            ".json",
            ".xls",
            ".xlsx",
            ".doc",
            ".docx",
            ".ppt",
            ".pptx",
            ".csv"
        }

    async def download_assets_from_html(
        self,
        base_url: str,
        html
    ):
        """
        Parse HTML and download all assets. Accepts either an
        HTML string or a BS4 soup/Tag so process_page can pass
        an already-parsed tree.
        """

        soup = _ensure_soup(html)

        asset_urls = set()

        # Images
        for img in soup.find_all("img", src=True):

            src = img["src"].strip()

            full_url = urljoin(base_url, src)

            asset_urls.add(full_url)

        # CSS
        for link in soup.find_all("link", href=True):

            href = link["href"].strip()

            rel = link.get("rel", [])

            if "stylesheet" in rel:
                asset_urls.add(
                    urljoin(base_url, href)
                )

        # JS
        for script in soup.find_all(
            "script",
            src=True
        ):

            src = script["src"].strip()

            asset_urls.add(
                urljoin(base_url, src)
            )

        logger.info(
            f"Found {len(asset_urls)} assets"
        )

        results = await asyncio.gather(
            *(self._download_safe(u) for u in asset_urls),
            return_exceptions=False,
        )

        downloaded_assets = [
            {"url": url, "local_path": str(path)}
            for url, path in results
            if path is not None
        ]

        return downloaded_assets

    async def _download_safe(self, asset_url):
        try:
            local_path = await self.download_asset(asset_url)
            return asset_url, local_path
        except Exception as e:
            logger.warning(
                f"Asset download failed {asset_url}: {e}"
            )
            return asset_url, None

    async def download_asset(
        self,
        asset_url: str
    ):
        """
        Download single asset.
        """

        asset_url = self.normalize_url(asset_url)

        cached = self.downloaded.get(asset_url)
        if cached is not None:
            logger.debug(f"Already downloaded: {asset_url}")
            return cached

        parsed = urlparse(asset_url)

        extension = Path(parsed.path).suffix.lower()

        if extension not in self.allowed_extensions:

            logger.debug(
                f"Skipped unsupported asset: "
                f"{asset_url}"
            )

            return None

        local_path = self.url_to_local_path(
            asset_url
        )

        # Skip re-downloading assets that are already on disk
        # from a prior run. Saves a HEAD+GET roundtrip per
        # cache-bustable asset on resume.
        if local_path.exists():
            self.downloaded[asset_url] = local_path
            return local_path

        local_path.parent.mkdir(
            parents=True,
            exist_ok=True
        )

        try:
            async with self.semaphore:
                response = await self.client.get(asset_url)
                response.raise_for_status()

                async with aiofiles.open(
                    local_path,
                    "wb"
                ) as f:
                    await f.write(response.content)

            self.downloaded[asset_url] = local_path

            logger.debug(
                f"Downloaded asset: {local_path}"
            )

            return local_path

        except Exception as e:

            logger.warning(
                f"Failed asset download "
                f"{asset_url}: {e}"
            )

            return None

    async def aclose(self):
        if self._owns_client:
            await self.client.aclose()

    def url_to_local_path(
        self,
        url: str
    ):
        """
        Convert asset URL -> local path.

        Example:

        https://api.intra-mart.jp/assets/logo.png

        =>
        output/assets/api.intra-mart.jp/assets/logo.png
        """

        parsed = urlparse(url)

        domain = parsed.netloc

        path = parsed.path.strip("/")

        if not path:

            filename = self.hash_url(url)

            path = f"unknown/{filename}"

        local_path = (
            self.output_dir /
            domain /
            path
        )

        return local_path

    def hash_url(
        self,
        url: str
    ):
        """
        Generate stable filename hash
        """

        return hashlib.md5(
            url.encode("utf-8")
        ).hexdigest()

    def normalize_url(
        self,
        url: str
    ):
        """
        Remove fragments and query strings.

        Query strings are stripped so that cache-busted URLs
        (e.g. logo.png?v=1 vs logo.png?v=2) deduplicate to the
        same local file path -- avoiding repeated downloads
        that overwrite the same file and avoiding stale dedup
        keys that point to a path written by a different URL.
        """

        parsed = urlparse(url)

        clean = parsed._replace(fragment="", query="")

        return clean.geturl()

    def rewrite_asset_links(
        self,
        html,
        base_url: str,
        current_output_path: Path
    ):
        """
        Rewrite HTML asset links to local filesystem paths.
        Accepts either an HTML string or a BS4 soup/Tag.
        Returns the rewritten HTML as a string.
        """

        soup = _ensure_soup(html)


        
        # Rewrite image src
        for img in soup.find_all("img", src=True):

            src = img["src"]

            full_url = urljoin(base_url, src)

            if Path(urlparse(full_url).path).suffix.lower() not in self.allowed_extensions:
                continue

            relative_path = self.make_relative_asset_path(
                full_url,
                current_output_path
            )

            img["src"] = relative_path

        # Rewrite CSS href -- ONLY rewrite tags we actually
        # download (rel=stylesheet). Other <link> tags
        # (canonical, alternate, icon, preload, ...) point to
        # files we never fetched, so rewriting their href
        # would produce dead local references.
        for link in soup.find_all(
            "link",
            href=True
        ):

            rel = link.get("rel", [])

            if "stylesheet" not in rel:
                continue

            href = link["href"]

            full_url = urljoin(base_url, href)

            if Path(urlparse(full_url).path).suffix.lower() not in self.allowed_extensions:
                continue

            relative_path = self.make_relative_asset_path(
                full_url,
                current_output_path
            )

            link["href"] = relative_path
            
        # Rewrite JS src
        for script in soup.find_all(
            "script",
            src=True
        ):

            src = script["src"]

            full_url = urljoin(base_url, src)

            if Path(urlparse(full_url).path).suffix.lower() not in self.allowed_extensions:
                continue

            relative_path = self.make_relative_asset_path(
                full_url,
                current_output_path
            )

            script["src"] = relative_path

        return str(soup)
    
    def make_relative_asset_path(
        self,
        asset_url: str,
        current_output_path: Path
    ):
        local_path = self.url_to_local_path(asset_url)

        relative = os.path.relpath(
            local_path,
            current_output_path.parent
        )

        return relative.replace("\\", "/")