adlist_mikrotik/source_conversion.py

"""Validate and convert complete upstream lists before returning HTTP 200.

The output is spooled to disk when it grows, avoiding a large in-memory
response. Validation and download errors happen before response headers are
sent to the browser or reverse proxy.
"""

import codecs
import hashlib
import re
import tempfile
from dataclasses import dataclass
from typing import BinaryIO, Callable


class InvalidSourceError(Exception):
    """An upstream body is not a supported, decodable hosts list."""


class SourceTooLargeError(InvalidSourceError):
    """The source or converted result exceeds a configured limit."""


@dataclass
class ConversionResult:
    output: BinaryIO
    lines_read: int
    lines_emitted: int
    output_size: int
    content_sha256: str


def _get_encoding(content_type: str) -> str:
    match = re.search(
        r'(?:^|;)\s*charset\s*=\s*(?:"([^"]+)"|([^;\s]+))',
        content_type or "", re.IGNORECASE,
    )
    encoding = (match.group(1) or match.group(2)) if match else "utf-8-sig"
    try:
        encoding = codecs.lookup(encoding).name
    except (LookupError, ValueError) as exc:
        raise InvalidSourceError("Unsupported source text encoding.") from exc
    # The reader splits on ASCII newline octets. Multi-byte encodings need a
    # different framing protocol and must not be interpreted as hosts text.
    if encoding in {"utf-16", "utf-16-le", "utf-16-be", "utf-32", "utf-32-le", "utf-32-be"}:
        raise InvalidSourceError("Unsupported source list encoding.")
    return encoding


def _bounded_lines(response, *, chunk_size: int, line_limit: int, source_limit: int):
    """Return raw lines, or None for overlong lines; never buffer a huge line."""
    pending = bytearray()
    skipping = False
    total = 0

    for chunk in response.iter_content(chunk_size=chunk_size):
        if not chunk:
            continue
        if not isinstance(chunk, bytes):
            raise InvalidSourceError("Source returned invalid binary data.")
        total += len(chunk)
        if total > source_limit:
            raise SourceTooLargeError("Source list exceeds the maximum allowed size.")
        pieces = chunk.split(b"\n")
        for part in pieces[:-1]:
            if skipping or len(pending) + len(part) > line_limit:
                yield None
            else:
                pending.extend(part)
                yield bytes(pending).rstrip(b"\r")
            pending.clear()
            skipping = False
        last = pieces[-1]
        if not skipping:
            if len(pending) + len(last) > line_limit:
                pending.clear()
                skipping = True
            else:
                pending.extend(last)

    if skipping:
        yield None
    elif pending:
        yield bytes(pending).rstrip(b"\r")


def convert_source(
    response,
    target_ip: str,
    convert_line: Callable[[str, str], str | None],
    *,
    read_chunk: int,
    line_limit: int,
    source_limit: int,
    output_limit: int,
    spool_memory_limit: int = 2 * 1024 * 1024,
    debug_lines: list[str] | None = None,
) -> ConversionResult:
    """Consume, decode and check the whole source before returning output.

    The caller owns and must close result.output. On failure, the spool is
    closed here, so no partial result can be published or cached as valid.
    """
    encoding = _get_encoding(response.headers.get("Content-Type", ""))
    length = response.headers.get("Content-Length", "")
    if length:
        try:
            if int(length) > source_limit:
                raise SourceTooLargeError("Source list exceeds the maximum allowed size.")
        except ValueError:
            pass

    output = tempfile.SpooledTemporaryFile(max_size=spool_memory_limit, mode="w+b")
    lines_read = 0
    lines_emitted = 0
    written = 0
    content_hasher = hashlib.sha256()

    def append(data: bytes):
        nonlocal written
        if written + len(data) > output_limit:
            raise SourceTooLargeError("Converted output exceeds the maximum allowed size.")
        output.write(data)
        written += len(data)
        content_hasher.update(data)

    try:
        if debug_lines is not None and debug_lines:
            append(("\n".join(debug_lines) + "\n").encode("utf-8"))

        for raw in _bounded_lines(
            response,
            chunk_size=read_chunk,
            line_limit=line_limit,
            source_limit=source_limit,
        ):
            lines_read += 1
            if raw is None:
                if debug_lines is not None and lines_read <= 5:
                    append(b"# [DEBUG] skipped an overlong line\n")
                continue
            if b"\x00" in raw:
                raise InvalidSourceError("Source contains binary data instead of a text list.")
            try:
                line = raw.decode(encoding, errors="strict")
                if lines_read == 1:
                    line = line.lstrip("\ufeff")
            except UnicodeError as exc:
                raise InvalidSourceError("Unable to decode source list text (character encoding).") from exc
            try:
                converted = convert_line(line, target_ip)
            except (TypeError, ValueError) as exc:
                raise InvalidSourceError("Unable to process this list format.") from exc
            if converted:
                lines_emitted += 1
                append((converted + "\n").encode("utf-8"))
            if debug_lines is not None and lines_read <= 5:
                preview = line[:200].replace("\r", "\\r").replace("\n", "\\n")
                append(f"# [DEBUG] preview {lines_read}: {preview}\n".encode("utf-8"))

        if lines_emitted == 0:
            raise InvalidSourceError(
                "No domains found in a supported format. Check the URL and list contents."
            )
        if debug_lines is not None:
            append(
                f"# [DEBUG] summary: read={lines_read}, emitted={lines_emitted}\n".encode("utf-8")
            )
        output.seek(0)
        return ConversionResult(output, lines_read, lines_emitted, written, content_hasher.hexdigest())
    except BaseException:
        output.close()
        raise


def iter_completed_bytes(output: BinaryIO, chunk_size: int):
    """The WSGI iterable must produce bytes, especially for WsgiToAsgi."""
    try:
        while True:
            data = output.read(chunk_size)
            if not data:
                break
            yield data
    finally:
        output.close()