From 018bd84e19a48240047b9e33626252d308d0268d Mon Sep 17 00:00:00 2001 From: Harsh Tita Date: Fri, 11 Sep 2026 14:58:44 -0700 Subject: [PATCH] fix(deploy-on-aws): resolve AttributeError from ET.Element type annotations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit defusedxml.ElementTree does not expose Element or ElementTree — it wraps only the parsing functions (parse, fromstring, iterparse) for security. Using ET.Element as a type annotation causes an AttributeError at import time on Python 3.10+ where annotations are evaluated eagerly, crashing all diagram post-processing scripts before any function is called. This affects all agent platforms (Claude Code, Codex, Cursor, Kiro) since the crash is in the Python scripts themselves, not the agent runtime. Fix: import Element and ElementTree from xml.etree.ElementTree for type annotations only. The defusedxml import is retained for all actual XML parsing calls (ET.parse, ET.fromstring, ET.iterparse) which is where the Bandit B314/B405 and Semgrep XXE security requirements apply. The Element and ElementTree classes carry no parsing security risk. Affected scripts: - scripts/lib/fix_step_badges.py - scripts/lib/fix_icon_colors.py - scripts/lib/fix_nesting.py - scripts/lib/post_process_drawio.py Validated: all 4 scripts import cleanly, fix_step_badges runs against example diagrams, and all 125 existing tests pass. By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of the project license. --- .../scripts/lib/fix_icon_colors.py | 5 +++-- .../deploy-on-aws/scripts/lib/fix_nesting.py | 11 ++++++----- .../scripts/lib/fix_step_badges.py | 17 +++++++++-------- .../scripts/lib/post_process_drawio.py | 9 +++++---- 4 files changed, 23 insertions(+), 19 deletions(-) diff --git a/plugins/deploy-on-aws/scripts/lib/fix_icon_colors.py b/plugins/deploy-on-aws/scripts/lib/fix_icon_colors.py index 913776b3..2b7ab4d0 100644 --- a/plugins/deploy-on-aws/scripts/lib/fix_icon_colors.py +++ b/plugins/deploy-on-aws/scripts/lib/fix_icon_colors.py @@ -10,6 +10,7 @@ import argparse import defusedxml.ElementTree as ET +from xml.etree.ElementTree import Element, ElementTree # Broken shape names → correct shape names SHAPE_RENAMES: dict[str, str] = { @@ -195,7 +196,7 @@ def _extract_color(value: str) -> str | None: return None -def fix_icon_colors(tree: ET.ElementTree, verbose: bool = False) -> int: +def fix_icon_colors(tree: ElementTree, verbose: bool = False) -> int: """Fix icon fillColor, container tint/stroke, and broken shape names. 1. Rename broken resIcon shapes (e.g., iam → identity_and_access_management) @@ -208,7 +209,7 @@ def fix_icon_colors(tree: ET.ElementTree, verbose: bool = False) -> int: fixed = 0 # Build a map of cell ID → cell element - cells: dict[str, ET.Element] = {} + cells: dict[str, Element] = {} for cell in root_elem.iter("mxCell"): cid = cell.get("id") if cid: diff --git a/plugins/deploy-on-aws/scripts/lib/fix_nesting.py b/plugins/deploy-on-aws/scripts/lib/fix_nesting.py index 9080e68e..eae68dae 100644 --- a/plugins/deploy-on-aws/scripts/lib/fix_nesting.py +++ b/plugins/deploy-on-aws/scripts/lib/fix_nesting.py @@ -17,6 +17,7 @@ import argparse import defusedxml.ElementTree as ET +from xml.etree.ElementTree import Element, ElementTree def get_style_dict(style_str: str) -> dict[str, str]: @@ -54,7 +55,7 @@ def set_style_value(style_str: str, key: str, value: str) -> str: return ";".join(parts) + ";" -def get_geometry(cell: ET.Element) -> tuple[float, float, float, float] | None: +def get_geometry(cell: Element) -> tuple[float, float, float, float] | None: for geom in cell: if geom.tag == "mxGeometry" and geom.get("as") == "geometry": x = float(geom.get("x", "0")) @@ -65,7 +66,7 @@ def get_geometry(cell: ET.Element) -> tuple[float, float, float, float] | None: return None -def offset_geometry(cell: ET.Element, dx: float, dy: float) -> None: +def offset_geometry(cell: Element, dx: float, dy: float) -> None: for geom in cell: if geom.tag == "mxGeometry" and geom.get("as") == "geometry": if geom.get("relative") == "1": @@ -77,7 +78,7 @@ def offset_geometry(cell: ET.Element, dx: float, dy: float) -> None: return -def is_region_container(cell: ET.Element) -> bool: +def is_region_container(cell: Element) -> bool: style = cell.get("style", "") style_dict = get_style_dict(style) return ( @@ -86,10 +87,10 @@ def is_region_container(cell: ET.Element) -> bool: ) -def fix_nesting(tree: ET.ElementTree, verbose: bool = False) -> int: +def fix_nesting(tree: ElementTree, verbose: bool = False) -> int: root_elem = tree.getroot() - cells: dict[str, ET.Element] = {} + cells: dict[str, Element] = {} for cell in root_elem.iter("mxCell"): cid = cell.get("id") if cid: diff --git a/plugins/deploy-on-aws/scripts/lib/fix_step_badges.py b/plugins/deploy-on-aws/scripts/lib/fix_step_badges.py index 0599a74d..41730a8f 100644 --- a/plugins/deploy-on-aws/scripts/lib/fix_step_badges.py +++ b/plugins/deploy-on-aws/scripts/lib/fix_step_badges.py @@ -21,6 +21,7 @@ import math import re import defusedxml.ElementTree as ET +from xml.etree.ElementTree import Element, ElementTree from dataclasses import dataclass @@ -88,7 +89,7 @@ def get_style_dict(style_str: str) -> dict[str, str]: return result -def get_geometry(cell: ET.Element) -> Rect | None: +def get_geometry(cell: Element) -> Rect | None: for geom in cell: if geom.tag == "mxGeometry" and geom.get("as") == "geometry": if geom.get("relative") == "1": @@ -102,8 +103,8 @@ def get_geometry(cell: ET.Element) -> Rect | None: def resolve_edge_label_position( - cell: ET.Element, - cells: dict[str, ET.Element], + cell: Element, + cells: dict[str, Element], geom_cache: dict[str, Rect], ) -> Rect | None: """Resolve an edge label's absolute position by finding the midpoint @@ -158,7 +159,7 @@ def resolve_edge_label_position( def resolve_absolute( cell_id: str, - cells: dict[str, ET.Element], + cells: dict[str, Element], geom_cache: dict[str, Rect], ) -> Rect | None: if cell_id in geom_cache: @@ -199,7 +200,7 @@ def resolve_absolute( return abs_rect -def is_on_diagram_badge(cell: ET.Element) -> bool: +def is_on_diagram_badge(cell: Element) -> bool: """On-diagram step badge: fillColor=#007CBD, numeric value, not in legend.""" style = get_style_dict(cell.get("style", "")) fill = style.get("fillColor", "").upper() @@ -219,7 +220,7 @@ def is_on_diagram_badge(cell: ET.Element) -> bool: return bool(re.match(r"^\d{1,2}$", stripped)) -def classify_cell(cell: ET.Element) -> str: +def classify_cell(cell: Element) -> str: """Classify a cell as 'badge', 'obstacle', or 'skip'.""" cell_id = cell.get("id", "") if cell_id in ("0", "1"): @@ -309,14 +310,14 @@ def compute_min_clearance( def fix_badges( - tree: ET.ElementTree, + tree: ElementTree, clearance: float = 10.0, verbose: bool = False, ) -> int: """Fix badge overlaps in-place. Returns number of badges moved.""" root_elem = tree.getroot() - cells: dict[str, ET.Element] = {} + cells: dict[str, Element] = {} for cell in root_elem.iter("mxCell"): cid = cell.get("id") if cid: diff --git a/plugins/deploy-on-aws/scripts/lib/post_process_drawio.py b/plugins/deploy-on-aws/scripts/lib/post_process_drawio.py index 3142ee2e..6fea7e05 100644 --- a/plugins/deploy-on-aws/scripts/lib/post_process_drawio.py +++ b/plugins/deploy-on-aws/scripts/lib/post_process_drawio.py @@ -16,6 +16,7 @@ import os import sys import defusedxml.ElementTree as ET +from xml.etree.ElementTree import Element, ElementTree from pathlib import Path MAX_FILE_SIZE = 2 * 1024 * 1024 # 2 MB @@ -55,7 +56,7 @@ def get_style_dict(style_str: str) -> dict[str, str]: return result -def get_geometry(cell: ET.Element) -> dict[str, float] | None: +def get_geometry(cell: Element) -> dict[str, float] | None: for geom in cell: if geom.tag == "mxGeometry" and geom.get("as") == "geometry": if geom.get("relative") == "1": @@ -69,7 +70,7 @@ def get_geometry(cell: ET.Element) -> dict[str, float] | None: return None -def set_geometry(cell: ET.Element, **kwargs: float) -> None: +def set_geometry(cell: Element, **kwargs: float) -> None: for geom in cell: if geom.tag == "mxGeometry" and geom.get("as") == "geometry": for k, v in kwargs.items(): @@ -78,7 +79,7 @@ def set_geometry(cell: ET.Element, **kwargs: float) -> None: return -def fix_placement(tree: ET.ElementTree, verbose: bool = False) -> int: +def fix_placement(tree: ElementTree, verbose: bool = False) -> int: """Move external actors outside the AWS Cloud boundary. External actors must be: @@ -218,7 +219,7 @@ def fix_placement(tree: ET.ElementTree, verbose: bool = False) -> int: return moved -def fix_legend_size(tree: ET.ElementTree, verbose: bool = False) -> int: +def fix_legend_size(tree: ElementTree, verbose: bool = False) -> int: """Resize legend panel to match the diagram's main content height. Finds the legend-outer group and the AWS Cloud / Region group,