Try-except and use get_lib_path

aikido_zen/vulnerabilities/sql_injection/__init__.py

@@ -2,41 +2,42 @@
 SQL Injection algorithm
 """
 
-import os
 import re
 import ctypes
 from aikido_zen.helpers.logging import logger
 from .map_dialect_to_rust_int import map_dialect_to_rust_int
-
-current_dir = os.path.dirname(os.path.abspath(__file__))
-lib_path = os.path.join(current_dir, "../../lib", "libzen_internals.so")
-internals_lib = ctypes.CDLL(lib_path)
+from .get_lib_path import get_binary_path
 
 
 def detect_sql_injection(query, user_input, dialect):
     """
     Execute this to check if the query is actually a SQL injection
     """
-    query_l = query.lower()
-    userinput_l = user_input.lower()
-    if should_return_early(query_l, userinput_l):
-        return False
-
-    query_bytes = query_l.encode("utf-8")
-    userinput_bytes = userinput_l.encode("utf-8")
-    dialect_int = map_dialect_to_rust_int(dialect)
-    c_int_res = internals_lib.detect_sql_injection(
-        query_bytes, userinput_bytes, dialect_int
-    )
+    try:
+        internals_lib = ctypes.CDLL(get_binary_path())
+        query_l = query.lower()
+        userinput_l = user_input.lower()
+        if should_return_early(query_l, userinput_l):
+            return False
 
-    # This means that an error occurred in the library
-    if c_int_res == 2:
-        logger.debug(
-            "Unable to check for SQL Injection, an error occurred in the library"
+        query_bytes = query_l.encode("utf-8")
+        userinput_bytes = userinput_l.encode("utf-8")
+        dialect_int = map_dialect_to_rust_int(dialect)
+        c_int_res = internals_lib.detect_sql_injection(
+            query_bytes, userinput_bytes, dialect_int
         )
-        return False
 
-    return bool(c_int_res)
+        # This means that an error occurred in the library
+        if c_int_res == 2:
+            logger.debug(
+                "Unable to check for SQL Injection, an error occurred in the library"
+            )
+            return False
+
+        return bool(c_int_res)
+    except Exception as e:
+        logger.debug("Exception in SQL algo: %s", e)
+    return False
 
 
 def should_return_early(query, user_input):

aikido_zen/vulnerabilities/sql_injection/get_lib_path.py

@@ -0,0 +1,32 @@
+"""Exports get_binary_path"""
+
+import platform
+import os
+
+
+def get_binary_path():
+    """Returns an absolute path for Rust binary file"""
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    lib_path = os.path.join(current_dir, "../../lib", get_file_name())
+    return lib_path
+
+
+def get_file_name():
+    """Gives you the file name for the binary based on platform info"""
+    os_name = platform.system().lower()
+    architecture = platform.architecture()[0].lower()
+    file_name = "libzen_internals_"
+
+    if "aarch64" in architecture:
+        file_name += "aarch64-"
+    elif "64" in architecture:
+        file_name += "x86_64-"
+
+    if os_name == "windows":
+        file_name += "pc-windows-gnu.dll"  # Windows
+    elif os_name == "darwin":
+        file_name += "apple-darwin.dylib"  # macOS
+    elif os_name in ["linux", "linux2"]:
+        file_name += "unknown-linux-gnu.so"  # Linux
+
+    return file_name