diff --git a/aikido_firewall/__init__.py b/aikido_firewall/__init__.py index 82bbb9ae..8eccc6e1 100644 --- a/aikido_firewall/__init__.py +++ b/aikido_firewall/__init__.py @@ -35,5 +35,6 @@ def protect(module="any", server=True): # Import sinks import aikido_firewall.sinks.pymysql + import aikido_firewall.sinks.mysqlclient logger.info("Aikido python firewall started") diff --git a/aikido_firewall/sinks/mysqlclient.py b/aikido_firewall/sinks/mysqlclient.py new file mode 100644 index 00000000..801df40f --- /dev/null +++ b/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"), "MySQLdb.connections.query", 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 diff --git a/sample-apps/django-mysql/manage.py b/sample-apps/django-mysql/manage.py index f66fda4b..15437c74 100755 --- a/sample-apps/django-mysql/manage.py +++ b/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