jb/diagram-design
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(import): harden draw.io extraction

Browse Source
This commit is contained in:
Cathryn Lavery
2026-08-11 13:55:37 -07:00
parent ced802198d
commit 6508d2e43e
7 changed files with 201 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
skills/diagram-design/scripts/drawio_extract.py
+82 -15
View File
@@ -39,6 +39,12 @@ from xml.etree import ElementTree as ET
# --------------------------------------------------------------------------
PNG_MAGIC = b"\x89PNG\r\n\x1a\n"
MAX_INPUT_BYTES = 32 * 1024 * 1024
MAX_XML_BYTES = 64 * 1024 * 1024
class PayloadTooLarge(ValueError):
"""Raised when compressed metadata expands beyond the supported limit."""
def _fail(msg: str) -> "NoReturn": # type: ignore[valid-type]
@@ -46,6 +52,37 @@ def _fail(msg: str) -> "NoReturn": # type: ignore[valid-type]
raise SystemExit(2)
def _reject_unsafe_xml(xml: str, source: str) -> None:
"""Reject declarations that can make XML parsing expand external data."""
upper = xml.upper()
if "<!DOCTYPE" in upper or "<!ENTITY" in upper:
_fail(f"{source}: DTD and entity declarations are not supported")
def _decompress_limited(data: bytes, wbits: int, limit: int = MAX_XML_BYTES) -> bytes:
"""Decompress without allowing a small payload to expand without bound."""
decompressor = zlib.decompressobj(wbits)
output = bytearray()
chunk = data
while chunk:
remaining = limit + 1 - len(output)
if remaining <= 0:
raise PayloadTooLarge(f"decoded payload exceeds {limit} bytes")
output.extend(decompressor.decompress(chunk, remaining))
if len(output) > limit:
raise PayloadTooLarge(f"decoded payload exceeds {limit} bytes")
chunk = decompressor.unconsumed_tail
if not decompressor.eof:
raise zlib.error("incomplete compressed payload")
remaining = limit + 1 - len(output)
if remaining <= 0:
raise PayloadTooLarge(f"decoded payload exceeds {limit} bytes")
output.extend(decompressor.flush(remaining))
if len(output) > limit:
raise PayloadTooLarge(f"decoded payload exceeds {limit} bytes")
return bytes(output)
def _inflate(payload: str) -> str | None:
"""Undo draw.io's base64 + raw-deflate + URL-encoding pipeline."""
try:
@@ -54,7 +91,11 @@ def _inflate(payload: str) -> str | None:
return None
for wbits in (-15, 15, 47):
try:
text = zlib.decompress(raw, wbits).decode("utf-8", "replace")
text = _decompress_limited(raw, wbits).decode("utf-8", "replace")
except PayloadTooLarge:
_fail(
f"decoded diagram exceeds the {MAX_XML_BYTES // (1024 * 1024)} MiB limit"
)
except Exception:
continue
# draw.io URL-encodes before deflating; unquote is a no-op if it didn't.
@@ -68,8 +109,12 @@ def _png_embedded_xml(data: bytes) -> str | None:
while pos + 8 <= len(data):
(length,) = struct.unpack(">I", data[pos : pos + 4])
ctype = data[pos + 4 : pos + 8]
body = data[pos + 8 : pos + 8 + length]
pos += 12 + length
body_end = pos + 8 + length
chunk_end = body_end + 4
if chunk_end > len(data):
_fail("PNG has a truncated metadata chunk")
body = data[pos + 8 : body_end]
pos = chunk_end
if ctype not in (b"tEXt", b"zTXt", b"iTXt"):
if ctype == b"IEND":
break
@@ -77,27 +122,44 @@ def _png_embedded_xml(data: bytes) -> str | None:
key, _, rest = body.partition(b"\x00")
if key.lower() != b"mxfile":
continue
if ctype == b"tEXt":
value = rest
elif ctype == b"zTXt":
value = zlib.decompress(rest[1:])
else: # iTXt: compression flag, method, lang, translated key, text
flag = rest[0:1]
tail = rest[2:].split(b"\x00", 2)[-1]
value = zlib.decompress(tail) if flag == b"\x01" else tail
try:
if ctype == b"tEXt":
value = rest
elif ctype == b"zTXt":
value = _decompress_limited(rest[1:], 15)
else: # iTXt: compression flag, method, lang, translated key, text
flag = rest[0:1]
tail = rest[2:].split(b"\x00", 2)[-1]
value = _decompress_limited(tail, 15) if flag == b"\x01" else tail
except PayloadTooLarge:
_fail(
f"embedded PNG diagram exceeds the "
f"{MAX_XML_BYTES // (1024 * 1024)} MiB limit"
)
except (IndexError, ValueError, zlib.error):
_fail("PNG has invalid compressed draw.io metadata")
return unquote(value.decode("utf-8", "replace"))
return None
def _svg_embedded_xml(text: str) -> str | None:
match = re.search(r'content\s*=\s*"((?:[^"\\]|\\.)*)"', text[:400000])
if not match:
return None
return html.unescape(match.group(1))
for match in re.finditer(
r"\bcontent\s*=\s*([\"'])(.*?)\1", text, flags=re.IGNORECASE | re.DOTALL
):
candidate = html.unescape(match.group(2))
if "<mxfile" in candidate or "<mxGraphModel" in candidate:
return candidate
return None
def load_mxfile(path: Path) -> str:
"""Return the ``<mxfile>`` (or bare ``<mxGraphModel>``) XML for any input."""
size = path.stat().st_size
if size > MAX_INPUT_BYTES:
_fail(
f"{path.name}: input is {size} bytes; maximum is "
f"{MAX_INPUT_BYTES // (1024 * 1024)} MiB"
)
data = path.read_bytes()
if data.startswith(PNG_MAGIC):
xml = _png_embedded_xml(data)
@@ -296,6 +358,7 @@ def parse_page(diagram: ET.Element, index: int) -> Page:
inflated = _inflate(text) if text else None
if not inflated:
return page
_reject_unsafe_xml(inflated, f"page {index}")
model = ET.fromstring(inflated)
if model.tag != "mxGraphModel":
found = model.find(".//mxGraphModel")
@@ -452,6 +515,7 @@ def parse_page(diagram: ET.Element, index: int) -> Page:
def parse_file(path: Path) -> list[Page]:
xml = load_mxfile(path)
_reject_unsafe_xml(xml, path.name)
try:
root = ET.fromstring(xml)
except ET.ParseError as exc:
@@ -766,6 +830,9 @@ def main(argv: list[str] | None = None) -> int:
parser.add_argument("--out", help="write to this path instead of stdout")
args = parser.parse_args(argv)
if args.max_rows < 1:
parser.error("--max-rows must be at least 1")
path = Path(args.file)
if not path.is_file():
_fail(f"{path}: no such file")