platform-ci/platform_ci/safe_xml.py

"""Parse untrusted XML without ever reading a DTD or an entity.

One expat pass refuses any DOCTYPE, entity declaration or external entity reference before the document is built, so nothing is
expanded or fetched, and the check holds in every encoding expat accepts (UTF-8 and UTF-16 alike), unlike a scan of the raw bytes.
"""

from xml.etree import ElementTree
from xml.parsers import expat


class ForbiddenDeclarationError(Exception):
   """Raised from inside the parser when a document declares a DTD or an entity."""


def parse_without_declarations(content: bytes) -> ElementTree.Element:
   """Build the element tree in one expat pass that refuses DTDs and entities.

   The public expat parser is used directly because the C-accelerated `ElementTree.XMLParser` exposes no handlers. Malformed XML
   raises `expat.ExpatError`; a declaration raises `ForbiddenDeclarationError`.
   """
   builder = ElementTree.TreeBuilder()
   parser = expat.ParserCreate()
   parser.StartDoctypeDeclHandler = _reject_declaration
   parser.EntityDeclHandler = _reject_declaration
   parser.ExternalEntityRefHandler = _reject_declaration
   parser.StartElementHandler = builder.start
   parser.EndElementHandler = builder.end
   parser.CharacterDataHandler = builder.data
   parser.Parse(content, True)
   return builder.close()


def _reject_declaration(*_arguments: object) -> int:
   """Refuse any DOCTYPE, ENTITY or external entity reference the parser reports."""
   raise ForbiddenDeclarationError