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 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


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("Nieznane kodowanie tekstu z serwera zrodlowego.") 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("Nieobslugiwane kodowanie pliku listy.")
    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("Zrodlo nie zwrocilo poprawnych danych binarnych.")
        total += len(chunk)
        if total > source_limit:
            raise SourceTooLargeError("Plik zrodlowy przekracza dopuszczalny rozmiar.")
        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("Plik zrodlowy przekracza dopuszczalny rozmiar.")
        except ValueError:
            pass

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

    def append(data: bytes):
        nonlocal written
        if written + len(data) > output_limit:
            raise SourceTooLargeError("Wynik konwersji przekracza dopuszczalny rozmiar.")
        output.write(data)
        written += len(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] pominieto zbyt dlugi wiersz\n")
                continue
            if b"\x00" in raw:
                raise InvalidSourceError("Zrodlo zawiera dane binarne zamiast listy tekstowej.")
            try:
                line = raw.decode(encoding, errors="strict")
                if lines_read == 1:
                    line = line.lstrip("\ufeff")
            except UnicodeError as exc:
                raise InvalidSourceError("Nie mozna zdekodowac tekstu listy (kodowanie).") from exc
            try:
                converted = convert_line(line, target_ip)
            except (TypeError, ValueError) as exc:
                raise InvalidSourceError("Nie mozna przetworzyc formatu listy.") 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] podglad {lines_read}: {preview}\n".encode("utf-8"))

        if lines_emitted == 0:
            raise InvalidSourceError(
                "Nie znaleziono domen w obslugiwanym formacie. Sprawdz adres i zawartosc listy."
            )
        if debug_lines is not None:
            append(
                f"# [DEBUG] podsumowanie: przeczytano={lines_read}, wyemitowano={lines_emitted}\n".encode("utf-8")
            )
        output.seek(0)
        return ConversionResult(output, lines_read, lines_emitted, written)
    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()