Merge pull request #103 from AikidoSec/AIK-3455

Add support for lxml

aikido_firewall/__init__.py

@@ -32,6 +32,7 @@ def protect(module="any", server=True):
     import aikido_firewall.sources.django
     import aikido_firewall.sources.flask
     import aikido_firewall.sources.xml
+    import aikido_firewall.sources.lxml
 
     import aikido_firewall.sources.gunicorn
     import aikido_firewall.sources.uwsgi

aikido_firewall/sources/lxml.py

@@ -0,0 +1,46 @@
+"""
+Sink module for `xml`, python's built-in function
+"""
+
+import copy
+import importhook
+from aikido_firewall.helpers.extract_data_from_xml_body import (
+    extract_data_from_xml_body,
+)
+from aikido_firewall.background_process.packages import add_wrapped_package
+
+
+@importhook.on_import("lxml.etree")
+def on_lxml_import(eltree):
+    """
+    Hook 'n wrap on `lxml.etree`.
+    - Wrap on fromstring() function
+    - Wrap on
+    Returns : Modified `lxml.etree` object
+    """
+    modified_eltree = importhook.copy_module(eltree)
+
+    former_fromstring = copy.deepcopy(eltree.fromstring)
+
+    def aikido_fromstring(text, *args, **kwargs):
+        res = former_fromstring(text, *args, **kwargs)
+        extract_data_from_xml_body(user_input=text, root_element=res)
+        return res
+
+    former_fromstringlist = copy.deepcopy(eltree.fromstringlist)
+
+    def aikido_fromstringlist(strings, *args, **kwargs):
+        res = former_fromstringlist(strings, *args, **kwargs)
+        for string in strings:
+            extract_data_from_xml_body(user_input=string, root_element=res)
+        return res
+
+    # pylint: disable=no-member
+    setattr(eltree, "fromstring", aikido_fromstring)
+    setattr(modified_eltree, "fromstring", aikido_fromstring)
+
+    # pylint: disable=no-member
+    setattr(eltree, "fromstringlist", aikido_fromstringlist)
+    setattr(modified_eltree, "fromstringlist", aikido_fromstringlist)
+    add_wrapped_package("lxml")
+    return modified_eltree

end2end/flask_postgres_xml_lxml_test.py

@@ -0,0 +1,28 @@
+import pytest
+import requests
+# 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"
+
+def test_safe_response_with_firewall():
+    xml_data = '<dogs><dog dog_name="Bobby" /></dogs>'
+    res = requests.post(post_url_fw, data=xml_data)
+    assert res.status_code == 200
+
+
+def test_safe_response_without_firewall():
+    xml_data = '<dogs><dog dog_name="Bobby" /></dogs>'
+    res = requests.post(post_url_nofw, data=xml_data)
+    assert res.status_code == 200
+
+
+def test_dangerous_response_with_firewall():
+    xml_data = '<dogs><dog dog_name="Malicious dog\', TRUE); -- " /></dogs>'
+    res = requests.post(post_url_fw, data=xml_data)
+    assert res.status_code == 500
+
+def test_dangerous_response_without_firewall():
+    xml_data = '<dogs><dog dog_name="Malicious dog\', TRUE); -- " /></dogs>'
+    res = requests.post(post_url_nofw, data=xml_data)
+    assert res.status_code == 200
+

sample-apps/flask-postgres-xml/app.py

@@ -10,6 +10,7 @@
 from flask import Flask, render_template, request
 import psycopg2
 import xml.etree.ElementTree as ET
+import lxml.etree as ET2
 
 app = Flask(__name__)
 if __name__ == '__main__':
@@ -61,6 +62,21 @@ def post_upload_xml():
 def post_xml():
     raw_xml = request.data.decode('utf-8')
     root = ET.fromstring(raw_xml)
+    ET2.fromstring(raw_xml)
+    conn = get_db_connection()
+    cursor =  conn.cursor()
+    for dog in root.findall('dog'):
+        dog_name = dog.get('dog_name')
+        cursor.execute(f"INSERT INTO dogs (dog_name, isAdmin) VALUES ('%s', FALSE)" % (dog_name))
+        conn.commit()
+    cursor.close()
+    conn.close()
+    return f'Dogs created successfully'
+
+@app.route("/xml_post_lxml", methods=['POST'])
+def post_xml():
+    raw_xml = request.data.decode('utf-8')
+    root = ET2.fromstring(raw_xml)
     conn = get_db_connection()
     cursor =  conn.cursor()
     for dog in root.findall('dog'):

sample-apps/flask-postgres-xml/requirements.txt

@@ -1,3 +1,4 @@
 flask==2.3.3
 psycopg2-binary
 cryptography
+lxml