Add wrapping for clickhouse-driver

README.md

@@ -52,6 +52,7 @@ Zen for Python 3 is compatible with:
 * ✅ [`psycopg`](https://pypi.org/project/psycopg) ^3.1
 * ✅ [`asyncpg`](https://pypi.org/project/asyncpg) ^0.27
 * ✅ [`motor`](https://pypi.org/project/motor/) (See `pymongo` version)
+* ✅ [`clickhouse-driver`](https://pypi.org/project/clickhouse-driver)
 
 ## Reporting to your Aikido Security dashboard
 

aikido_zen/__init__.py

@@ -51,6 +51,8 @@ def protect(mode="daemon"):
     import aikido_zen.sinks.psycopg2
     import aikido_zen.sinks.psycopg
     import aikido_zen.sinks.asyncpg
+    import aikido_zen.sinks.clickhouse_driver
+
     import aikido_zen.sinks.builtins
     import aikido_zen.sinks.os
     import aikido_zen.sinks.shutil

aikido_zen/sinks/clickhouse_driver.py

@@ -0,0 +1,24 @@
+from aikido_zen.helpers.get_argument import get_argument
+from aikido_zen.sinks import before, on_import, patch_function
+from aikido_zen.vulnerabilities import run_vulnerability_scan
+
+
+@before
+def _execute(func, instance, args, kwargs):
+    kind = "sql_injection"
+    op = "clickhouse_driver.Client.execute"
+    query = get_argument(args, kwargs, 0, "query")
+    run_vulnerability_scan(kind, op, args=(query, "clickhouse"))
+
+
+@on_import("clickhouse_driver", package="clickhouse_driver")
+def patch(m):
+    """
+    patching module clickhouse_driver
+    - patches clickhouse_driver.Client.execute
+    - patches clickhouse_driver.Client.execute_iter
+    - patches clickhouse_driver.Client.execute_with_progress
+    """
+    patch_function(m, "Client.execute", _execute)
+    patch_function(m, "Client.execute_iter", _execute)
+    patch_function(m, "Client.execute_with_progress", _execute)

aikido_zen/sinks/tests/clickhouse_driver_test.py

@@ -0,0 +1,130 @@
+import aikido_zen.sinks.clickhouse_driver
+import pytest
+from aikido_zen.background_process import reset_comms
+from aikido_zen.context import Context
+from aikido_zen.errors import AikidoSQLInjection
+
+
+class Context1(Context):
+    def __init__(self, body):
+        self.cookies = {}
+        self.headers = {}
+        self.remote_address = "1.1.1.1"
+        self.method = "POST"
+        self.url = "url"
+        self.query = {}
+        self.body = body
+        self.source = "express"
+        self.route = "/"
+        self.parsed_userinput = {}
+
+
+@pytest.fixture(autouse=True)
+def set_blocking_to_true(monkeypatch):
+    monkeypatch.setenv("AIKIDO_BLOCK", "1")
+
+
+@pytest.fixture
+def client():
+    from clickhouse_driver import Client
+
+    return Client(
+        host="127.0.0.1", port=9000, user="default", password="", database="default"
+    )
+
+
+def test_client_execute_without_context(client):
+    reset_comms()
+    dog_name = "Steve"
+    sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
+    client.execute(sql)
+
+
+def test_client_execute_safe(client):
+    reset_comms()
+    dog_name = "Steve"
+    sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
+    Context1({"dog_name": dog_name}).set_as_current_context()
+    client.execute(sql)
+
+
+def test_client_execute_unsafe(client, monkeypatch):
+    reset_comms()
+    dog_name = "Malicious dog', 1); -- "
+    sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
+    Context1({"dog_name": dog_name}).set_as_current_context()
+
+    with pytest.raises(AikidoSQLInjection):
+        client.execute(sql)
+
+    monkeypatch.setenv("AIKIDO_BLOCK", "0")
+    client.execute(sql)
+
+
+def test_cursor_execute_safe():
+    from clickhouse_driver import connect
+
+    conn = connect("clickhouse://localhost:9000")
+    reset_comms()
+    dog_name = "Steve"
+    sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
+    Context1({"dog_name": dog_name}).set_as_current_context()
+    conn.cursor().execute(sql)
+
+
+def test_cursor_execute_unsafe(monkeypatch):
+    from clickhouse_driver import connect
+
+    conn = connect("clickhouse://localhost:9000")
+    reset_comms()
+    dog_name = "Malicious dog', 1); -- "
+    sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
+    Context1({"dog_name": dog_name}).set_as_current_context()
+
+    with pytest.raises(AikidoSQLInjection):
+        conn.cursor().execute(sql)
+
+    monkeypatch.setenv("AIKIDO_BLOCK", "0")
+    conn.cursor().execute(sql)
+
+
+def test_client_execute_with_progress_safe(client):
+    reset_comms()
+    dog_name = "Steve"
+    sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
+    Context1({"dog_name": dog_name}).set_as_current_context()
+    client.execute_with_progress(sql)
+
+
+def test_client_execute_with_progress_unsafe(client, monkeypatch):
+    reset_comms()
+    dog_name = "Malicious dog', 1); -- "
+    sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
+    Context1({"dog_name": dog_name}).set_as_current_context()
+
+    with pytest.raises(AikidoSQLInjection):
+        client.execute_with_progress(sql)
+
+    monkeypatch.setenv("AIKIDO_BLOCK", "0")
+    client.execute_with_progress(sql)
+
+
+def test_client_execute_iter_safe(client):
+    reset_comms()
+    dog_name = "Steve"
+    sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
+    Context1({"dog_name": dog_name}).set_as_current_context()
+    client.execute_iter(sql)
+
+
+def test_client_execute_iter_unsafe(client, monkeypatch):
+    reset_comms()
+    dog_name = "Malicious dog', 1); -- "
+    sql = "INSERT INTO dogs (dog_name, isAdmin) VALUES ('{}' , 0)".format(dog_name)
+    Context1({"dog_name": dog_name}).set_as_current_context()
+
+    with pytest.raises(AikidoSQLInjection):
+        client.execute_iter(sql)
+
+    monkeypatch.setenv("AIKIDO_BLOCK", "0")
+    client.execute_iter(sql)

aikido_zen/vulnerabilities/sql_injection/map_dialect_to_rust_int.py

@@ -4,6 +4,7 @@
 
 DIALECTS = {
     "generic": 0,
+    "clickhouse": 3,
     "mysql": 8,
     "postgres": 9,
     "sqlite": 12,