refactor: expand map_dialect_to_rust_int fn

to its natural habitat

Author: isimluk

aikido_zen/vulnerabilities/sql_injection/map_dialect_to_rust_int.py

@@ -5,7 +5,6 @@
 import ctypes
 import threading
 from .get_lib_path import get_binary_path
-from .map_dialect_to_rust_int import map_dialect_to_rust_int
 from ...helpers.encode_safely import encode_safely
 
 class __Singleton(type):
@@ -19,6 +18,15 @@ def __call__(cls, *args, **kwargs):
         return cls._instances[cls]
 
 class ZenInternal(metaclass=__Singleton):
+    # Reference : [rust lib]/src/sql_injection/helpers/select_dialect_based_on_enum.rs
+    SQL_DIALECTS = {
+        "generic": 0,
+        "clickhouse": 3,
+        "mysql": 8,
+        "postgres": 9,
+        "sqlite": 12,
+    }
+
     def __init__(self):
         self._lib = ctypes.CDLL(get_binary_path())
         self._lib.detect_sql_injection.argtypes = [
@@ -37,7 +45,7 @@ def detect_sql_injection(self, query, user_input, dialect):
         userinput_buffer = (ctypes.c_uint8 * len(userinput_bytes)).from_buffer_copy(
             userinput_bytes
         )
-        dialect_int = map_dialect_to_rust_int(dialect)
+        dialect_int = self.SQL_DIALECTS[dialect]
         return self._lib.detect_sql_injection(
             query_buffer,
             len(query_bytes),