Add mysqlclient support

aikido_firewall/__init__.py

@@ -22,9 +22,9 @@
     if server:
         start_background_process()
     else:
         logger.debug("Not starting background process")
     if module == "server-only":
         return
 
     # Import sources
     if module == "django":
@@ -35,5 +35,6 @@
 
     # Import sinks
     import aikido_firewall.sinks.pymysql
+    import aikido_firewall.sinks.mysqlclient
 
     logger.info("Aikido python firewall started")

aikido_firewall/sinks/mysqlclient.py

@@ -0,0 +1,45 @@
+"""
+Sink module for `mysqlclient`
+"""
+
+import copy
+import json
+from importlib.metadata import version
+import importhook
+from aikido_firewall.context import get_current_context
+from aikido_firewall.vulnerabilities.sql_injection.check_context_for_sql_injection import (
+    check_context_for_sql_injection,
+)
+from aikido_firewall.vulnerabilities.sql_injection.dialects import MySQL
+from aikido_firewall.helpers.logging import logger
+
+
+@importhook.on_import("MySQLdb.connections")
+def on_mysqlclient_import(mysql):
+    """
+    Hook 'n wrap on `MySQLdb.connections`
+    Our goal is to wrap the query() function of the Connection class :
+    https://github.com/PyMySQL/mysqlclient/blob/9fd238b9e3105dcbed2b009a916828a38d1f0904/src/MySQLdb/connections.py#L257
+    Returns : Modified MySQLdb.connections object
+    """
+    modified_mysql = importhook.copy_module(mysql)
+
+    prev_query_function = copy.deepcopy(mysql.Connection.query)
+
+    def aikido_new_query(_self, sql):
+        logger.debug("Wrapper - `mysqlclient` version : %s", version("mysqlclient"))
+
+        context = get_current_context()
+        result = check_context_for_sql_injection(
+            sql.decode("utf-8"), "Test_op", context, MySQL()
+        )
+
+        logger.debug("sql_injection results : %s", json.dumps(result))
+        if result:
+            raise Exception("SQL Injection [aikido_firewall]")
+        return prev_query_function(_self, sql)
+
+    # pylint: disable=no-member
+    setattr(mysql.Connection, "query", aikido_new_query)
+    logger.debug("Wrapped `mysqlclient` module")
+    return modified_mysql

aikido_firewall/vulnerabilities/sql_injection/init_test.py

@@ -0,0 +1,315 @@
+import os
+import pytest
+from aikido_firewall.vulnerabilities.sql_injection import detect_sql_injection
+from aikido_firewall.vulnerabilities.sql_injection.dialects import MySQL
+from aikido_firewall.vulnerabilities.sql_injection.dialects import Postgres
+from aikido_firewall.vulnerabilities.sql_injection.consts import SQL_DANGEROUS_IN_STRING
+
+BAD_SQL_COMMANDS = [
+    "Roses are red insErt are blue",
+    "Roses are red cREATE are blue",
+    "Roses are red drop are blue",
+    "Roses are red updatE are blue",
+    "Roses are red SELECT are blue",
+    "Roses are red dataBASE are blue",
+    "Roses are red alter are blue",
+    "Roses are red grant are blue",
+    "Roses are red savepoint are blue",
+    "Roses are red commit are blue",
+    "Roses are red or blue",
+    "Roses are red and lovely",
+    "This is a group_concat_test",
+    "I'm writting you",
+    "Termin;ate",
+    "Roses <> violets",
+    "Roses < Violets",
+    "Roses > Violets",
+    "Roses != Violets",
+]
+
+GOOD_SQL_COMMANDS = [
+    "Roses are red rollbacks are blue",
+    "Roses are red truncates are blue",
+    "Roses are reddelete are blue",
+    "Roses are red WHEREis blue",
+    "Roses are red ORis isAND",
+    "abcdefghijklmnop@hotmail.com",
+    "steve@yahoo.com",
+    "I was benchmark ing",
+    "We were delay ed",
+    "I will waitfor you",
+    "#",
+    "'",
+]
+
+IS_NOT_INJECTION = [
+    ["'UNION 123' UNION \"UNION 123\"", "UNION 123"],
+    ["'union'  is not \"UNION\"", "UNION!"],
+    ['"UNION;"', "UNION;"],
+    ["SELECT * FROM table", "*"],
+    ['"COPY/*"', "COPY/*"],
+    ["'union'  is not \"UNION--\"", "UNION--"],
+]
+
+IS_INJECTION = [
+    ["'union'  is not UNION", "UNION"],
+    ["UNTER;", "UNTER;"],
+]
+
+
+def is_sql_injection(sql, input):
+    result = detect_sql_injection(sql, input, MySQL())
+    assert result == True, f"Expected SQL injection for SQL: {sql} and input: {input}"
+    result = detect_sql_injection(sql, input, Postgres())
+    assert result == True, f"Expected SQL injection for SQL: {sql} and input: {input}"
+    return result
+
+
+def is_not_sql_injection(sql, input):
+    result = detect_sql_injection(sql, input, MySQL())
+    assert (
+        result == False
+    ), f"Expected no SQL injection for SQL: {sql} and input: {input}"
+    result = detect_sql_injection(sql, input, Postgres())
+    assert (
+        result == False
+    ), f"Expected no SQL injection for SQL: {sql} and input: {input}"
+    return result
+
+
+def test_bad_sql_commands():
+    for sql in BAD_SQL_COMMANDS:
+        is_sql_injection(sql, sql)
+
+
+def test_good_sql_commands():
+    for sql in GOOD_SQL_COMMANDS:
+        is_not_sql_injection(sql, sql)
+
+
+def test_is_injection():
+    for sql, input in IS_INJECTION:
+        is_sql_injection(sql, input)
+
+
+def test_is_not_injection():
+    for sql, input in IS_NOT_INJECTION:
+        is_not_sql_injection(sql, input)
+
+
+def test_allow_escape_sequences():
+    is_sql_injection("SELECT * FROM users WHERE id = 'users\\'", "users\\")
+    is_sql_injection("SELECT * FROM users WHERE id = 'users\\\\'", "users\\\\")
+    is_not_sql_injection("SELECT * FROM users WHERE id = '\nusers'", "\nusers")
+    is_not_sql_injection("SELECT * FROM users WHERE id = '\rusers'", "\rusers")
+    is_not_sql_injection("SELECT * FROM users WHERE id = '\tusers'", "\tusers")
+
+
+def test_user_input_inside_in():
+    is_sql_injection("SELECT * FROM users WHERE id IN ('123')", "'123'")
+    is_not_sql_injection("SELECT * FROM users WHERE id IN (123)", "123")
+    is_not_sql_injection("SELECT * FROM users WHERE id IN (123, 456)", "123")
+    is_not_sql_injection("SELECT * FROM users WHERE id IN (123, 456)", "456")
+    is_not_sql_injection("SELECT * FROM users WHERE id IN ('123')", "123")
+    is_not_sql_injection("SELECT * FROM users WHERE id IN (13,14,15)", "13,14,15")
+    is_not_sql_injection("SELECT * FROM users WHERE id IN (13, 14, 154)", "13, 14, 154")
+    is_sql_injection(
+        "SELECT * FROM users WHERE id IN (13, 14, 154) OR (1=1)", "13, 14, 154) OR (1=1"
+    )
+
+
+def test_check_string_safely_escaped():
+    is_sql_injection(
+        "SELECT * FROM comments WHERE comment = 'I'm writting you'", "I'm writting you"
+    )
+    is_sql_injection(
+        'SELECT * FROM comments WHERE comment = "I"m writting you"', 'I"m writting you'
+    )
+    is_sql_injection("SELECT * FROM `comm`ents`", "`comm`ents")
+    is_not_sql_injection(
+        'SELECT * FROM comments WHERE comment = "I\'m writting you"', "I'm writting you"
+    )
+    is_not_sql_injection(
+        "SELECT * FROM comments WHERE comment = 'I\"m writting you'", 'I"m writting you'
+    )
+    is_not_sql_injection(
+        'SELECT * FROM comments WHERE comment = "I\`m writting you"', "I`m writting you"
+    )
+    is_not_sql_injection("SELECT * FROM `comm'ents`", "comm'ents")
+
+
+def test_not_flag_select_queries():
+    is_not_sql_injection("SELECT * FROM users WHERE id = 1", "SELECT")
+
+
+def test_not_flag_escaped_hash():
+    is_not_sql_injection("SELECT * FROM hashtags WHERE name = '#hashtag'", "#hashtag")
+
+
+def test_comment_same_as_user_input():
+    is_sql_injection(
+        "SELECT * FROM hashtags WHERE name = '-- Query by name' -- Query by name",
+        "-- Query by name",
+    )
+
+
+def test_input_occurs_in_comment():
+    is_not_sql_injection(
+        "SELECT * FROM hashtags WHERE name = 'name' -- Query by name", "name"
+    )
+
+
+def test_user_input_is_multiline():
+    is_sql_injection("SELECT * FROM users WHERE id = 'a'\nOR 1=1#'", "a'\nOR 1=1#")
+    is_not_sql_injection("SELECT * FROM users WHERE id = 'a\nb\nc';", "a\nb\nc")
+
+
+def test_user_input_is_longer_than_query():
+    is_not_sql_injection("SELECT * FROM users", "SELECT * FROM users WHERE id = 'a'")
+
+
+def test_multiline_queries():
+    is_sql_injection(
+        """
+        SELECT * FROM `users`
+        WHERE id = 123
+        """,
+        "users`",
+    )
+
+    is_sql_injection(
+        """
+        SELECT *
+        FROM users
+        WHERE id = '1' OR 1=1
+        """,
+        "1' OR 1=1",
+    )
+
+    is_sql_injection(
+        """
+        SELECT *
+        FROM users
+        WHERE id = '1' OR 1=1
+        AND is_escaped = '1'' OR 1=1'
+        """,
+        "1' OR 1=1",
+    )
+
+    is_sql_injection(
+        """
+        SELECT *
+        FROM users
+        WHERE id = '1' OR 1=1
+        AND is_escaped = "1' OR 1=1"
+        """,
+        "1' OR 1=1",
+    )
+
+    is_not_sql_injection(
+        """
+        SELECT * FROM `users`
+        WHERE id = 123
+        """,
+        "123",
+    )
+
+    is_not_sql_injection(
+        """
+        SELECT * FROM `us``ers`
+        WHERE id = 123
+        """,
+        "users",
+    )
+
+    is_not_sql_injection(
+        """
+        SELECT * FROM users
+        WHERE id = 123
+        """,
+        "123",
+    )
+
+    is_not_sql_injection(
+        """
+        SELECT * FROM users
+        WHERE id = '123'
+        """,
+        "123",
+    )
+
+    is_not_sql_injection(
+        """
+        SELECT *
+        FROM users
+        WHERE is_escaped = "1' OR 1=1"
+        """,
+        "1' OR 1=1",
+    )
+
+
+def test_lowercased_input_sql_injection():
+    sql = """
+        SELECT id,
+               email,
+               password_hash,
+               registered_at,
+               is_confirmed,
+               first_name,
+               last_name
+        FROM users WHERE email_lowercase = '' or 1=1 -- a'
+    """
+    expected_sql_injection = "' OR 1=1 -- a"
+
+    assert is_sql_injection(sql, expected_sql_injection)
+
+
+@pytest.mark.parametrize("dangerous", SQL_DANGEROUS_IN_STRING)
+def test_dangerous_strings(dangerous):
+    input = f"{dangerous} a"
+    is_sql_injection(f"SELECT * FROM users WHERE {input}", input)
+
+
+def test_function_calls_as_sql_injections():
+    is_sql_injection("foobar()", "foobar()")
+    is_sql_injection("foobar(1234567)", "foobar(1234567)")
+    is_sql_injection("foobar       ()", "foobar       ()")
+    is_sql_injection(".foobar()", ".foobar()")
+    is_sql_injection("20+foobar()", "20+foobar()")
+    is_sql_injection("20-foobar(", "20-foobar(")
+    is_sql_injection("20<foobar()", "20<foobar()")
+    is_sql_injection("20*foobar  ()", "20*foobar  ()")
+    is_sql_injection("!foobar()", "!foobar()")
+    is_sql_injection("=foobar()", "=foobar()")
+    is_sql_injection("1foobar()", "1foobar()")
+    is_sql_injection("1foo_bar()", "1foo_bar()")
+    is_sql_injection("1foo-bar()", "1foo-bar()")
+    is_sql_injection("#foobar()", "#foobar()")
+
+    is_not_sql_injection("foobar)", "foobar)")
+    is_not_sql_injection("foobar      )", "foobar      )")
+    is_not_sql_injection("€foobar()", "€foobar()")
+
+
+def file_paths():
+    script_dir = os.path.dirname(__file__)
+    return [
+        # Taken from https://github.com/payloadbox/sql-injection-payload-list/tree/master
+        os.path.join(script_dir, "payloads/Auth_Bypass.txt"),
+        os.path.join(script_dir, "payloads/postgres.txt"),
+        os.path.join(script_dir, "payloads/mysql.txt"),
+        os.path.join(script_dir, "payloads/mssql_and_db2.txt"),
+    ]
+
+
+# Define the Pytest tests
+for file_path in file_paths():
+    with open(file_path, "r", encoding="utf-8") as file:
+        for line in file:
+            sql = line.rstrip("\n")
+
+            def test_sql_injection():
+                assert is_sql_injection(sql, sql)
+
+            def test_sql_injection_query():
+                assert is_sql_injection(f"SELECT * FROM users WHERE id = {sql}", sql)

aikido_firewall/vulnerabilities/sql_injection/uinput_occ_safely_encapsulated.py

@@ -44,7 +44,9 @@ def uinput_occ_safely_encapsulated(query, user_input):
     """
     This function will check if user input is actually just safely encapsulated in the query
     """
-    segments_in_between = get_current_and_next_segments(query.split(user_input))
+    segments_in_between = get_current_and_next_segments(
+        query.lower().split(user_input.lower())
+    )
 
     for segment in segments_in_between:
         current_seg, next_seg = segment

sample-apps/django-mysql/manage.py

@@ -1,14 +1,13 @@
 #!/usr/bin/env python
 """Django's command-line utility for administrative tasks."""
-import aikido_firewall # Aikido module
-aikido_firewall.protect("django")
-
 import os
 import sys
+import aikido_firewall # Aikido module
 
 
 def main():
     """Run administrative tasks."""
+    aikido_firewall.protect("django")
     os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'sample-django-mysql-app.settings')
     try:
         from django.core.management import execute_from_command_line