Scan XML element text and attributes at any depth

aikido_zen/helpers/extract_data_from_xml_body.py

@@ -5,7 +5,7 @@
 
 
 def extract_data_from_xml_body(user_input, root_element):
-    """Extracts all attributes from the xml and adds them to context"""
+    """Extracts all attributes and text from the xml and adds them to context"""
     try:
         context = ctx.get_current_context()
         if not context or not isinstance(context.body, str):
@@ -16,12 +16,19 @@ def extract_data_from_xml_body(user_input, root_element):
         if user_input != context.body:
             return
 
-        extracted_xml_attrs = context.xml
-        for el in root_element:
-            for k, v in el.items():
-                if not extracted_xml_attrs.get(k):
-                    extracted_xml_attrs[k] = set()
-                extracted_xml_attrs[k].add(v)
+        extracted_xml = context.xml
+        for element in root_element.iter():
+            if not isinstance(element.tag, str):
+                continue
+            for key, value in element.items():
+                extracted_xml.setdefault(key, set()).add(value)
+            for text in (element.text, element.tail):
+                if not text:
+                    continue
+                stripped = text.strip()
+                if not stripped:
+                    continue
+                extracted_xml.setdefault(element.tag, set()).update((text, stripped))
         context.set_as_current_context()
     except Exception as e:
         logger.debug("Exception occurred when extracting XML: %s", e)

aikido_zen/helpers/extract_data_from_xml_body_test.py

@@ -1,5 +1,6 @@
 import pytest
 from unittest.mock import MagicMock, patch
+import xml.etree.ElementTree as ET
 import aikido_zen.context as ctx
 from .extract_data_from_xml_body import (
     extract_data_from_xml_body,
@@ -17,21 +18,19 @@ def mock_context():
 def test_does_not_crash_when_context_none(mock_context):
     with patch("aikido_zen.context.get_current_context", return_value=None):
         user_input = "valid_input"
-        root_element = [
-            {"attr1": "value1", "attr2": "value2"},
-            {"attr1": "value3", "attr3": "value4"},
-        ]
+        root_element = ET.fromstring(
+            '<root><a attr1="value1" attr2="value2"/><b attr1="value3" attr3="value4"/></root>'
+        )
 
         extract_data_from_xml_body(user_input, root_element)
 
 
 def test_extract_data_from_xml_body_valid_input(mock_context):
     with patch("aikido_zen.context.get_current_context", return_value=mock_context):
         user_input = "valid_input"
-        root_element = [
-            {"attr1": "value1", "attr2": "value2"},
-            {"attr1": "value3", "attr3": "value4"},
-        ]
+        root_element = ET.fromstring(
+            '<root><a attr1="value1" attr2="value2"/><b attr1="value3" attr3="value4"/></root>'
+        )
 
         extract_data_from_xml_body(user_input, root_element)
 
@@ -45,7 +44,7 @@ def test_extract_data_from_xml_body_valid_input(mock_context):
 def test_extract_data_from_xml_body_invalid_user_input(mock_context):
     with patch("aikido_zen.context.get_current_context", return_value=mock_context):
         user_input = "invalid_input"
-        root_element = [{"attr1": "value1"}]
+        root_element = ET.fromstring('<root><a attr1="value1"/></root>')
 
         extract_data_from_xml_body(user_input, root_element)
 
@@ -55,7 +54,7 @@ def test_extract_data_from_xml_body_invalid_user_input(mock_context):
 def test_extract_data_from_xml_body_empty_root_element(mock_context):
     with patch("aikido_zen.context.get_current_context", return_value=mock_context):
         user_input = "valid_input"
-        root_element = []
+        root_element = ET.fromstring("<root></root>")
 
         extract_data_from_xml_body(user_input, root_element)
 
@@ -66,7 +65,7 @@ def test_extract_data_from_xml_body_non_string_context_body(mock_context):
     with patch("aikido_zen.context.get_current_context", return_value=mock_context):
         mock_context.body = 123  # Set body to a non-string value
         user_input = "valid_input"
-        root_element = [{"attr1": "value1"}]
+        root_element = ET.fromstring('<root><a attr1="value1"/></root>')
 
         extract_data_from_xml_body(user_input, root_element)
 
@@ -76,8 +75,8 @@ def test_extract_data_from_xml_body_non_string_context_body(mock_context):
 def test_extract_data_from_xml_body_multiple_calls(mock_context):
     with patch("aikido_zen.context.get_current_context", return_value=mock_context):
         user_input = "valid_input"
-        root_element1 = [{"attr1": "value1"}]
-        root_element2 = [{"attr1": "value2"}]
+        root_element1 = ET.fromstring('<root><a attr1="value1"/></root>')
+        root_element2 = ET.fromstring('<root><a attr1="value2"/></root>')
 
         extract_data_from_xml_body(user_input, root_element1)
         extract_data_from_xml_body(user_input, root_element2)
@@ -88,11 +87,9 @@ def test_extract_data_from_xml_body_multiple_calls(mock_context):
 def test_extract_data_from_xml_body_duplicate_attributes(mock_context):
     with patch("aikido_zen.context.get_current_context", return_value=mock_context):
         user_input = "valid_input"
-        root_element = [
-            {"attr1": "value1"},
-            {"attr1": "value1"},  # Duplicate
-            {"attr2": "value2"},
-        ]
+        root_element = ET.fromstring(
+            '<root><a attr1="value1"/><b attr1="value1"/><c attr2="value2"/></root>'
+        )
 
         extract_data_from_xml_body(user_input, root_element)
 
@@ -102,17 +99,91 @@ def test_extract_data_from_xml_body_duplicate_attributes(mock_context):
 def test_extract_data_from_xml_body_no_attributes(mock_context):
     with patch("aikido_zen.context.get_current_context", return_value=mock_context):
         user_input = "valid_input"
-        root_element = [{}]
+        root_element = ET.fromstring("<root><a/></root>")
 
         extract_data_from_xml_body(user_input, root_element)
 
         assert mock_context.xml == {}
 
 
+def test_extract_data_from_xml_body_element_text(mock_context):
+    with patch("aikido_zen.context.get_current_context", return_value=mock_context):
+        user_input = "valid_input"
+        root_element = ET.fromstring("<root><id>value1</id></root>")
+
+        extract_data_from_xml_body(user_input, root_element)
+
+        assert mock_context.xml == {"id": {"value1"}}
+
+
+def test_extract_data_from_xml_body_indented_element_text(mock_context):
+    with patch("aikido_zen.context.get_current_context", return_value=mock_context):
+        user_input = "valid_input"
+        root_element = ET.fromstring("<root><id>\n    value1\n  </id></root>")
+
+        extract_data_from_xml_body(user_input, root_element)
+
+        assert mock_context.xml == {"id": {"\n    value1\n  ", "value1"}}
+
+
+def test_extract_data_from_xml_body_text_between_elements(mock_context):
+    with patch("aikido_zen.context.get_current_context", return_value=mock_context):
+        user_input = "valid_input"
+        root_element = ET.fromstring("<root>before<a/>after</root>")
+
+        extract_data_from_xml_body(user_input, root_element)
+
+        assert mock_context.xml == {"root": {"before"}, "a": {"after"}}
+
+
+def test_extract_data_from_xml_body_root_attributes(mock_context):
+    with patch("aikido_zen.context.get_current_context", return_value=mock_context):
+        user_input = "valid_input"
+        root_element = ET.fromstring('<root attr1="value1"><a/></root>')
+
+        extract_data_from_xml_body(user_input, root_element)
+
+        assert mock_context.xml == {"attr1": {"value1"}}
+
+
+def test_extract_data_from_xml_body_nested_attributes(mock_context):
+    with patch("aikido_zen.context.get_current_context", return_value=mock_context):
+        user_input = "valid_input"
+        root_element = ET.fromstring('<root><a><b attr1="value1"/></a></root>')
+
+        extract_data_from_xml_body(user_input, root_element)
+
+        assert mock_context.xml == {"attr1": {"value1"}}
+
+
+@pytest.mark.parametrize(
+    "root_element", [None, "not an element", [{"attr1": "value1"}]]
+)
+def test_extract_data_from_xml_body_invalid_root_element(mock_context, root_element):
+    with patch("aikido_zen.context.get_current_context", return_value=mock_context):
+        extract_data_from_xml_body("valid_input", root_element)
+
+        assert mock_context.xml == {}
+
+
+def test_extract_data_from_xml_body_comments_and_processing_instructions(mock_context):
+    with patch("aikido_zen.context.get_current_context", return_value=mock_context):
+        parser = ET.XMLParser(
+            target=ET.TreeBuilder(insert_comments=True, insert_pis=True)
+        )
+        root_element = ET.fromstring(
+            "<root><!-- comment --><?pi data?><a>value1</a></root>", parser
+        )
+
+        extract_data_from_xml_body("valid_input", root_element)
+
+        assert mock_context.xml == {"a": {"value1"}}
+
+
 def test_extract_data_from_xml_body_context_set_as_current(mock_context):
     with patch("aikido_zen.context.get_current_context", return_value=mock_context):
         user_input = "valid_input"
-        root_element = [{"attr1": "value1"}]
+        root_element = ET.fromstring('<root><a attr1="value1"/></root>')
 
         extract_data_from_xml_body(user_input, root_element)
 
@@ -122,7 +193,9 @@ def test_extract_data_from_xml_body_context_set_as_current(mock_context):
 def test_extract_data_from_xml_body_bytes_user_input(mock_context):
     with patch("aikido_zen.context.get_current_context", return_value=mock_context):
         user_input = b"valid_input"
-        root_element = [{"attr1": "value1"}, {"attr2": "value2"}]
+        root_element = ET.fromstring(
+            '<root><a attr1="value1"/><b attr2="value2"/></root>'
+        )
 
         extract_data_from_xml_body(user_input, root_element)
 
@@ -132,7 +205,7 @@ def test_extract_data_from_xml_body_bytes_user_input(mock_context):
 def test_extract_data_from_xml_body_bytes_invalid_user_input(mock_context):
     with patch("aikido_zen.context.get_current_context", return_value=mock_context):
         user_input = b"invalid_input"
-        root_element = [{"attr1": "value1"}]
+        root_element = ET.fromstring('<root><a attr1="value1"/></root>')
 
         extract_data_from_xml_body(user_input, root_element)
 
@@ -144,20 +217,21 @@ def test_extract_data_from_xml_body_bytes_invalid_user_input(mock_context):
 )
 def test_extract_data_from_xml_body_byteslike_user_input(mock_context, user_input):
     with patch("aikido_zen.context.get_current_context", return_value=mock_context):
-        root_element = [{"attr1": "value1"}]
+        root_element = ET.fromstring('<root><a attr1="value1"/></root>')
 
         extract_data_from_xml_body(user_input, root_element)
 
         assert mock_context.xml == {"attr1": {"value1"}}
 
 
 def test_extract_data_from_xml_body_bytes_invalid_utf8_matches_body():
-    body = b'<xml><item id="1\xff"></item></xml>'
+    body = b'<xml><item id="1\xff">1\xff</item></xml>'
     context = ctx.Context(body=body, source="test")
+    root_element = ET.fromstring(context.body)
     context.set_as_current_context()
     try:
-        extract_data_from_xml_body(body, [{"id": "1\ufffd"}])
+        extract_data_from_xml_body(body, root_element)
 
-        assert context.xml == {"id": {"1\ufffd"}}
+        assert context.xml == {"id": {"1�"}, "item": {"1�"}}
     finally:
         ctx.current_context.set(None)

aikido_zen/sources/xml_sources/lxml_test.py

@@ -68,9 +68,12 @@ def test_parse_xml_with_set_context():
     context = get_current_context()
     assert context.xml == {
         "attr": {"chill"},
+        "name": {"Test Name"},
+        "ok": {"boomer"},
         "smth": {"2"},
         "smthelse": {"2"},
-        "test": {"test1"},
+        "test": {"test1", "test2"},
+        "value": {"42"},
     }
 
 
@@ -84,7 +87,10 @@ def test_parse_xml_bytes_with_set_context():
     context = get_current_context()
     assert context.xml == {
         "attr": {"chill"},
+        "name": {"Test Name"},
+        "ok": {"boomer"},
         "smth": {"2"},
         "smthelse": {"2"},
-        "test": {"test1"},
+        "test": {"test1", "test2"},
+        "value": {"42"},
     }

aikido_zen/sources/xml_sources/xml_test.py

@@ -70,9 +70,12 @@ def test_parse_xml_with_set_context():
     context = get_current_context()
     assert context.xml == {
         "attr": {"chill"},
+        "name": {"Test Name"},
+        "ok": {"boomer"},
         "smth": {"2"},
         "smthelse": {"2"},
-        "test": {"test1"},
+        "test": {"test1", "test2"},
+        "value": {"42"},
     }
 
 
@@ -86,7 +89,10 @@ def test_parse_xml_bytes_with_set_context():
     context = get_current_context()
     assert context.xml == {
         "attr": {"chill"},
+        "name": {"Test Name"},
+        "ok": {"boomer"},
         "smth": {"2"},
         "smthelse": {"2"},
-        "test": {"test1"},
+        "test": {"test1", "test2"},
+        "value": {"42"},
     }

end2end/flask_postgres_xml_lxml_test.py

@@ -6,6 +6,8 @@
 # e2e tests for flask_postgres sample app
 post_url_fw = "http://localhost:8092/xml_post_lxml"
 post_url_nofw = "http://localhost:8093/xml_post_lxml"
+text_url_fw = "http://localhost:8092/xml_post_text_lxml"
+text_url_nofw = "http://localhost:8093/xml_post_text_lxml"
 
 def test_safe_response_with_firewall():
     xml_data = '<dogs><dog dog_name="Bobby" /></dogs>'
@@ -25,8 +27,8 @@ def test_dangerous_response_with_firewall():
 
     time.sleep(5) # Wait for attack to be reported
     events = fetch_events_from_mock("http://localhost:5000")
-    attacks = filter_on_event_type(events, "detected_attack")
-    
+    attacks = [a for a in filter_on_event_type(events, "detected_attack") if a["request"]["route"] == "/xml_post_lxml"]
+
     assert len(attacks) == 1
     del attacks[0]["attack"]["stack"]
 
@@ -49,3 +51,33 @@ def test_dangerous_response_without_firewall():
     res = requests.post(post_url_nofw, data=xml_data)
     assert res.status_code == 200
 
+def test_dangerous_element_text_with_firewall():
+    xml_data = '<dogs><dog>Malicious dog\', TRUE); --</dog></dogs>'
+    res = requests.post(text_url_fw, data=xml_data)
+    assert res.status_code == 500
+
+    time.sleep(5) # Wait for attack to be reported
+    events = fetch_events_from_mock("http://localhost:5000")
+    attacks = [a for a in filter_on_event_type(events, "detected_attack") if a["request"]["route"] == "/xml_post_text_lxml"]
+
+    assert len(attacks) == 1
+    del attacks[0]["attack"]["stack"]
+
+    assert attacks[0]["attack"] == {
+        "blocked": True,
+        "kind": "sql_injection",
+        'metadata': {
+            'dialect': "postgres",
+            'sql': "INSERT INTO dogs (dog_name, isAdmin) VALUES ('Malicious dog', TRUE); --', FALSE)"
+        },
+        'operation': "psycopg2.Connection.Cursor.execute",
+        'pathToPayload': ".dog.[0]",
+        'payload': "\"Malicious dog', TRUE); --\"",
+        'source': "xml",
+        'user': None
+    }
+
+def test_dangerous_element_text_without_firewall():
+    xml_data = '<dogs><dog>Malicious dog\', TRUE); --</dog></dogs>'
+    res = requests.post(text_url_nofw, data=xml_data)
+    assert res.status_code == 200