# Copyright (C) 2026 Wilhelmus Dingenus de Pooter
# SPDX-License-Identifier: GPL-3.0-or-later
"""
enc_xml_safe.py — hardened XML parsing for every XML document the plugin reads.

The plugin parses XML from two kinds of sources: files it BUNDLES (feature
catalogues, symbol/style definitions) and files that arrive WITH CHART DATA
(exchange-set catalogues, Feature Catalogues a user installs). The second kind
is untrusted input, so ALL parsing goes through this one helper, which:

* rejects any document carrying a DTD or entity declaration outright. stdlib
  ElementTree expands *internal* entities, so a crafted DTD could mount a
  "billion laughs" entity-expansion DoS; no hydrographic XML format the plugin
  reads legitimately uses a DTD, so a blanket reject is both safe and correct.
  (External entities are never resolved by ElementTree, so there is no
  file-disclosure/XXE exposure either way — this closes the remaining vector.)
* parses the pre-screened bytes via the ElementTree instance API.

Pure stdlib — the QGIS Python environment does not ship defusedxml.
"""
import io
import xml.etree.ElementTree as ET


def _screen(data):
    """Reject DTD/entity declarations anywhere in the document (case-blind,
    UTF-16 aware). Returns the raw bytes unchanged when clean."""
    probe = data
    if probe[:2] in (b"\xff\xfe", b"\xfe\xff"):          # UTF-16 BOM
        probe = probe.decode("utf-16", "replace").encode("utf-8", "replace")
    up = probe.upper()
    if b"<!DOCTYPE" in up or b"<!ENTITY" in up:
        raise ValueError(
            "XML document contains a DTD/entity declaration, which is not "
            "permitted (rejected to prevent entity-expansion attacks).")
    return data


def parse_tree(source):
    """Defensively parse an XML file path or binary file object.

    Returns an ``xml.etree.ElementTree.ElementTree`` (so ``.getroot()`` /
    ``.write()`` work as usual). Raises ``ValueError`` for documents carrying
    DTD/entity declarations; propagates ``FileNotFoundError`` / ``ParseError``
    like the stdlib would.
    """
    if hasattr(source, "read"):
        data = source.read()
    else:
        with open(source, "rb") as fh:
            data = fh.read()
    if isinstance(data, str):
        data = data.encode("utf-8", "replace")
    tree = ET.ElementTree()
    tree.parse(io.BytesIO(_screen(data)))
    return tree


def parse_root(source):
    """Root element of a defensively parsed XML document (see parse_tree)."""
    return parse_tree(source).getroot()


def from_string(text):
    """Defensive replacement for the stdlib fromstring helper — parse XML held
    in a str or bytes object, returning the root element."""
    if isinstance(text, str):
        text = text.encode("utf-8", "replace")
    tree = ET.ElementTree()
    tree.parse(io.BytesIO(_screen(text)))
    return tree.getroot()
