update-lists.py

#!/usr/bin/env python3

import argparse
import asyncio
import base64
import hashlib
import json
import logging
import re
import shutil
import tempfile
from pathlib import Path

import aiohttp
import sentry_sdk
from aiohttp_retry import ExponentialRetry, RetryClient

logger = logging.getLogger("update_lists")
logging.basicConfig(
    level=logging.INFO,
    format=(
        "[%(asctime)s] {%(filename)s:%(lineno)d} %(funcName)s - "
        "%(levelname)s - %(message)s"
    ),
)

MAX_FILE_SIZE = 100 * 1024 * 1024  # 100 MB

sentry_sdk.init(enable_tracing=False)


def parse_arguments():
    parser = argparse.ArgumentParser(
        description="Tool to download the lists for an Adblock catalog"
    )
    parser.add_argument(
        "--adblock-catalog",
        type=str,
        help="the URL of the Adblock catalog",
        default="https://raw.githubusercontent.com/brave/adblock-resources/master/filter_lists/list_catalog.json",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        help="the directory to save the downloaded lists",
        default="lists",
    )
    return parser.parse_args()


def validate_checksum(filename):
    """Validate the checksum header"""
    data = Path(filename).read_bytes().decode("utf-8")

    # Extract and remove checksum
    checksum_pattern = re.compile(
        r"^\s*!\s*checksum[\s\-:]+([\w\+\/=]+).*\n", re.MULTILINE | re.IGNORECASE
    )
    match = checksum_pattern.search(data)
    if not match:
        logger.warning(f"Couldn't find a checksum in {filename}")
        return

    checksum = match.group(1)
    data = checksum_pattern.sub("", data, 1)

    # Normalize data
    data = re.sub(r"\r", "", data)
    data = re.sub(r"\n+\Z", "\n", data)

    # Calculate new checksum
    checksum_expected = hashlib.md5(data.encode("utf-8")).digest()
    checksum_expected = base64.b64encode(checksum_expected).decode().rstrip("=")

    # Compare checksums
    if checksum == checksum_expected:
        logging.info(f"Checksum is valid: {filename}")
    else:
        raise Exception(
            f"Wrong checksum, found {checksum}, "
            f"expected [{checksum_expected}] in {filename}"
        )


def move_downloaded_file(filename, url, output_dir):
    """
    Moves the downloaded file to the appropriate location in the output directory.

    Args:
        filename (str): The name of the downloaded file.
        url (str): The URL from which the file was downloaded.
        output_dir (str): The directory where the file should be moved.

    Returns:
        str: The path of the moved file.

    Notes:
        The filename is generated by hashing the URL.
    """
    output_file_name = hashlib.md5(url.encode("utf-8")).hexdigest() + ".txt"
    output_file_path = Path(output_dir) / output_file_name

    try:
        if url.startswith("https://easylist-downloads.adblockplus.org/"):
            print(f"ignoring checksum calculation for {url}")
        else:
            validate_checksum(filename)
        logger.info(f"moving {filename} to {output_file_path}")
        shutil.move(filename, output_file_path)
    except Exception:
        logger.exception(
            f"An exception happened while processing {filename} from {url}"
        )
        try:
            Path(filename).unlink()
        except Exception:
            logger.debug(
                "Failed to remove temporary file %s during error handling.",
                filename,
                exc_info=True,
            )

    return output_file_path


async def fetch_and_save_url(session, url, output_dir):
    """Fetch a URL and save it to a file, with retries for transient errors."""
    try:
        async with session.get(
            url, raise_for_status=True, timeout=aiohttp.ClientTimeout(total=60)
        ) as response:
            # Check if the response is successful
            if response.status == 200:
                # Create a temporary file
                with tempfile.NamedTemporaryFile(delete=False) as temp_file:
                    file_size = 0
                    # Write the response content to the temporary file
                    while True:
                        chunk = await response.content.read(1024)
                        if not chunk:
                            break

                        file_size += len(chunk)
                        if file_size > MAX_FILE_SIZE:
                            logger.error(
                                "Download failed for %s: file size (%d bytes) "
                                "exceeds the limit of %d bytes.",
                                url,
                                file_size,
                                MAX_FILE_SIZE,
                            )
                            Path(temp_file.name).unlink()
                            return

                        temp_file.write(chunk)
                    logger.info(f"downloaded {url}")
                    temp_file.close()
                    move_downloaded_file(temp_file.name, url, output_dir)
                return  # Success
            else:
                logger.error(
                    "Request to %s returned status %s, expected 200. Skipping.",
                    url,
                    response.status,
                )
                return
    except (aiohttp.ClientError, asyncio.TimeoutError):
        logger.exception(f"An exception occurred while processing {url}")
        return


async def main():
    args = parse_arguments()

    # Fetch the adblock catalog using aiohttp
    retry_options = ExponentialRetry(attempts=3)
    connector = aiohttp.TCPConnector(limit_per_host=3)

    async with RetryClient(
        connector=connector, retry_options=retry_options, connector_owner=False
    ) as session:
        async with session.get(
            args.adblock_catalog,
            timeout=aiohttp.ClientTimeout(total=60),
            raise_for_status=True,
        ) as response:
            adblock_catalog = await response.json(content_type="text/plain")

    adblock_lists = []
    metadata = {}
    for al in adblock_catalog:
        for src in al["sources"]:
            url = src["url"]
            adblock_lists.append(url)
            metadata[hashlib.md5(url.encode("utf-8")).hexdigest()] = url

    metadata_file = Path(args.output_dir) / "metadata.json"
    metadata_file.write_text(json.dumps(metadata, indent=4) + "\n")

    async with RetryClient(connector=connector, retry_options=retry_options) as session:
        tasks = [
            fetch_and_save_url(session, url, args.output_dir) for url in adblock_lists
        ]
        await asyncio.gather(*tasks)


if __name__ == "__main__":
    asyncio.run(main())