AIK-3168 : Pour sql_injection vulnerability code over to python firewall (including testing)

.pylintrc

@@ -4,5 +4,4 @@ disable = unused-argument
 ignore-patterns=(.)*_test\.py,test_(.)*\.py
 
 [MESSAGES CONTROL]
-# Disable the no-member warning [Currently enabled]
-#disable = no-member
+disable=too-few-public-methods

aikido_firewall/helpers/__init__.py

@@ -0,0 +1,8 @@
+"""
+init file for the helpers/ folder, imports helper function for easier accessibility
+"""
+
+from aikido_firewall.helpers.escape_string_regexp import escape_string_regexp
+from aikido_firewall.helpers.get_current_and_next_segments import (
+    get_current_and_next_segments,
+)

aikido_firewall/helpers/build_path_to_payload.py

@@ -0,0 +1,23 @@
+"""
+Helper function file, see funtion definition
+"""
+
+
+def build_path_to_payload(path_to_payload):
+    """
+    Create a string so people see the path to where
+    the injection took place
+    """
+    if len(path_to_payload) == 0:
+        return "."
+
+    result = ""
+    for part in path_to_payload:
+        if part["type"] == "object":
+            result += "." + part["key"]
+        elif part["type"] == "array":
+            result += f".[{part['index']}]"
+        elif part["type"] == "jwt":
+            result += "<jwt>"
+
+    return result

aikido_firewall/helpers/build_path_to_payload_test.py

@@ -0,0 +1,41 @@
+import pytest
+from aikido_firewall.helpers.build_path_to_payload import build_path_to_payload
+
+
+def test_build_path_to_payload_empty():
+    assert build_path_to_payload([]) == "."
+
+
+def test_build_path_to_payload_single_object():
+    path = [{"type": "object", "key": "name"}]
+    assert build_path_to_payload(path) == ".name"
+
+
+def test_build_path_to_payload_single_array():
+    path = [{"type": "array", "index": 0}]
+    assert build_path_to_payload(path) == ".[0]"
+
+
+def test_build_path_to_payload_single_jwt():
+    path = [{"type": "jwt"}]
+    assert build_path_to_payload(path) == "<jwt>"
+
+
+def test_build_path_to_payload_mixed_types():
+    path = [
+        {"type": "object", "key": "user"},
+        {"type": "array", "index": 2},
+        {"type": "jwt"},
+        {"type": "object", "key": "details"},
+        {"type": "array", "index": 1},
+    ]
+    assert build_path_to_payload(path) == ".user.[2]<jwt>.details.[1]"
+
+
+def test_build_path_to_payload_multiple_objects():
+    path = [
+        {"type": "object", "key": "user"},
+        {"type": "object", "key": "details"},
+        {"type": "object", "key": "address"},
+    ]
+    assert build_path_to_payload(path) == ".user.details.address"

aikido_firewall/helpers/escape_string_regexp.py

@@ -0,0 +1,17 @@
+"""
+Helper function file, see funtion definition
+"""
+
+import re
+
+
+def escape_string_regexp(string):
+    """
+    Escape characters with special meaning either inside or outside character sets.
+    Use a simple backslash escape when it's always valid, and a '\\xnn' escape
+    when the simpler form would be disallowed by Unicode patterns' stricter grammar.
+    Taken from https://github.com/sindresorhus/escape-string-regexp/
+    """
+    pattern = re.compile(r"[|\\{}()[\]^$+*?.]")
+    replace1 = re.sub(pattern, r"\\\g<0>", string)
+    return re.sub(r"-", r"\\x2d", replace1)

aikido_firewall/helpers/escape_string_regexp_test.py

@@ -0,0 +1,18 @@
+import re
+import pytest
+from aikido_firewall.helpers import escape_string_regexp
+
+
+def test_escape_string_regexp():
+    assert (
+        escape_string_regexp("\\ ^ $ * + ? . ( ) | { } [ ]")
+        == "\\\\ \\^ \\$ \\* \\+ \\? \\. \\( \\) \\| \\{ \\} \\[ \\]"
+    )
+
+
+def test_escape_hyphen_pcre():
+    assert escape_string_regexp("foo - bar") == "foo \\x2d bar"
+
+
+def test_escape_hyphen_unicode_flag():
+    assert re.match(escape_string_regexp("-"), "-u") is not None

aikido_firewall/helpers/extract_strings_from_user_input.py

@@ -0,0 +1,42 @@
+"""
+Helper function file, see funtion definition
+"""
+
+from aikido_firewall.helpers.try_decode_as_jwt import try_decode_as_jwt
+from aikido_firewall.helpers.build_path_to_payload import build_path_to_payload
+
+
+def extract_strings_from_user_input(obj, path_to_payload=None):
+    """
+    Extracts strings from an object (user input)
+    """
+    if path_to_payload is None:
+        path_to_payload = []
+
+    results = {}
+
+    if isinstance(obj, dict):
+        for key, value in obj.items():
+            results[key] = build_path_to_payload(path_to_payload)
+            for k, v in extract_strings_from_user_input(
+                value, path_to_payload + [{"type": "object", "key": key}]
+            ).items():
+                results[k] = v
+
+    if isinstance(obj, list):
+        for i, value in enumerate(obj):
+            for k, v in extract_strings_from_user_input(
+                value, path_to_payload + [{"type": "array", "index": i}]
+            ).items():
+                results[k] = v
+
+    if isinstance(obj, str):
+        results[obj] = build_path_to_payload(path_to_payload)
+        jwt = try_decode_as_jwt(obj)
+        if jwt[0]:
+            for k, v in extract_strings_from_user_input(
+                jwt[1], path_to_payload + [{"type": "jwt"}]
+            ).items():
+                results[k] = v
+
+    return results

aikido_firewall/helpers/extract_strings_from_user_input_test.py

@@ -0,0 +1,96 @@
+import pytest
+from aikido_firewall.helpers.extract_strings_from_user_input import (
+    extract_strings_from_user_input,
+)
+
+
+def from_obj(obj):
+    return dict(obj)
+
+
+def test_empty_object_returns_empty_dict():
+    assert extract_strings_from_user_input({}) == from_obj({})
+
+
+def test_extract_query_objects():
+    assert extract_strings_from_user_input({"age": {"$gt": "21"}}) == from_obj(
+        {"age": ".", "$gt": ".age", "21": ".age.$gt"}
+    )
+    assert extract_strings_from_user_input({"title": {"$ne": "null"}}) == from_obj(
+        {"title": ".", "$ne": ".title", "null": ".title.$ne"}
+    )
+    assert extract_strings_from_user_input(
+        {"age": "whaat", "user_input": ["whaat", "dangerous"]}
+    ) == from_obj(
+        {
+            "user_input": ".",
+            "age": ".",
+            "whaat": ".user_input.[0]",
+            "dangerous": ".user_input.[1]",
+        }
+    )
+
+
+def test_extract_cookie_objects():
+    assert extract_strings_from_user_input(
+        {"session": "ABC", "session2": "DEF"}
+    ) == from_obj(
+        {"session2": ".", "session": ".", "ABC": ".session", "DEF": ".session2"}
+    )
+    assert extract_strings_from_user_input(
+        {"session": "ABC", "session2": 1234}
+    ) == from_obj({"session2": ".", "session": ".", "ABC": ".session"})
+
+
+def test_extract_header_objects():
+    assert extract_strings_from_user_input(
+        {"Content-Type": "application/json"}
+    ) == from_obj({"Content-Type": ".", "application/json": ".Content-Type"})
+    assert extract_strings_from_user_input({"Content-Type": 54321}) == from_obj(
+        {"Content-Type": "."}
+    )
+    assert extract_strings_from_user_input(
+        {"Content-Type": "application/json", "ExtraHeader": "value"}
+    ) == from_obj(
+        {
+            "Content-Type": ".",
+            "application/json": ".Content-Type",
+            "ExtraHeader": ".",
+            "value": ".ExtraHeader",
+        }
+    )
+
+
+def test_extract_body_objects():
+    assert extract_strings_from_user_input(
+        {"nested": {"nested": {"$ne": None}}}
+    ) == from_obj({"nested": ".nested", "$ne": ".nested.nested"})
+    assert extract_strings_from_user_input(
+        {"age": {"$gt": "21", "$lt": "100"}}
+    ) == from_obj(
+        {"age": ".", "$lt": ".age", "$gt": ".age", "21": ".age.$gt", "100": ".age.$lt"}
+    )
+
+
+def test_decodes_jwts():
+    assert extract_strings_from_user_input(
+        {
+            "token": "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwidXNlcm5hbWUiOnsiJG5lIjpudWxsfSwiaWF0IjoxNTE2MjM5MDIyfQ._jhGJw9WzB6gHKPSozTFHDo9NOHs3CNOlvJ8rWy6VrQ"
+        }
+    ) == from_obj(
+        {
+            "token": ".",
+            "iat": ".token<jwt>",
+            "username": ".token<jwt>",
+            "sub": ".token<jwt>",
+            "1234567890": ".token<jwt>.sub",
+            "$ne": ".token<jwt>.username",
+            "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwidXNlcm5hbWUiOnsiJG5lIjpudWxsfSwiaWF0IjoxNTE2MjM5MDIyfQ._jhGJw9WzB6gHKPSozTFHDo9NOHs3CNOlvJ8rWy6VrQ": ".token",
+        }
+    )
+
+
+def test_jwt_as_string():
+    assert extract_strings_from_user_input(
+        {"header": "/;ping%20localhost;.e30=."}
+    ) == from_obj({"header": ".", "/;ping%20localhost;.e30=.": ".header"})

aikido_firewall/helpers/get_current_and_next_segments.py

@@ -0,0 +1,8 @@
+"""
+Helper function file, see function definition for more info
+"""
+
+
+def get_current_and_next_segments(array):
+    """Get the current and next segments of an array"""
+    return [(array[i], array[i + 1]) for i in range(len(array) - 1)]

aikido_firewall/helpers/get_current_and_next_segments_test.py

@@ -0,0 +1,28 @@
+import pytest
+from aikido_firewall.helpers.get_current_and_next_segments import (
+    get_current_and_next_segments,
+)
+
+
+def test_empty_array():
+    assert get_current_and_next_segments([]) == []
+
+
+def test_single_item_array():
+    assert get_current_and_next_segments(["a"]) == []
+
+
+def test_two_item_array():
+    assert get_current_and_next_segments(["a", "b"]) == [("a", "b")]
+
+
+def test_three_item_array():
+    assert get_current_and_next_segments(["a", "b", "c"]) == [("a", "b"), ("b", "c")]
+
+
+def test_four_item_array():
+    assert get_current_and_next_segments(["a", "b", "c", "d"]) == [
+        ("a", "b"),
+        ("b", "c"),
+        ("c", "d"),
+    ]

aikido_firewall/helpers/is_plain_object.py

@@ -0,0 +1,11 @@
+"""
+Helper function file, see funtion definition
+"""
+
+
+def is_plain_object(o):
+    """
+    Checks if the object is a plain object,
+    i.e. a dictionary in Python
+    """
+    return str(type(o)) == "<class 'dict'>"

aikido_firewall/helpers/is_plain_object_test.py

@@ -0,0 +1,34 @@
+import pytest
+from aikido_firewall.helpers.is_plain_object import is_plain_object
+
+
+def test_is_plain_object_true():
+    assert is_plain_object({}) == True
+    assert is_plain_object({"foo": "bar"}) == True
+
+
+def test_is_plain_object_false():
+    class Foo:
+        def __init__(self):
+            self.abc = {}
+
+    foo_instance = Foo()
+
+    assert is_plain_object("/foo/") == False
+    assert is_plain_object(lambda: None) == False
+    assert is_plain_object(1) == False
+    assert is_plain_object(["foo", "bar"]) == False
+    assert is_plain_object([]) == False
+    assert is_plain_object(foo_instance) == False
+    assert is_plain_object(None) == False
+
+
+def test_is_plain_object_modified_prototype():
+    class CustomConstructor:
+        pass
+
+    CustomConstructor.prototype = list
+
+    instance = CustomConstructor()
+
+    assert is_plain_object(instance) == False

aikido_firewall/helpers/try_decode_as_jwt.py

@@ -0,0 +1,25 @@
+"""
+Helper function file, see function docstrin
+"""
+
+import base64
+import json
+
+
+def try_decode_as_jwt(jwt):
+    """
+    This will try to decode a string as a JWT
+    """
+    if "." not in jwt:
+        return (False, None)
+
+    parts = jwt.split(".")
+
+    if len(parts) != 3:
+        return (False, None)
+
+    try:
+        jwt = json.loads(base64.b64decode(str.encode(parts[1]) + b"==").decode("utf-8"))
+        return (True, jwt)
+    except Exception:
+        return (False, None)