Use ctypes for entropy calculation (partially reverts redballoonsec…

…urity#150) (redballoonsecurity#483)

ofrak_core/CHANGELOG.md

@@ -34,6 +34,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - Change `FreeSpaceModifier` & `PartialFreeSpaceModifier` behavior: an optional stub that isn't free space can be provided and fill-bytes for free space can be specified. ([#409](https://github.com/redballoonsecurity/ofrak/pull/409))
 - `Resource.flush_to_disk` method renamed to `Resource.flush_data_to_disk`. ([#373](https://github.com/redballoonsecurity/ofrak/pull/373))
 - `build_image.py` supports building Docker images with OFRAK packages from any ancestor directory. ([#425](https://github.com/redballoonsecurity/ofrak/pull/425))
+- Partially reverted [#150](https://github.com/redballoonsecurity/ofrak/pull/150) so entropy C code is called with `ctypes` again, but maintaining the current API and automatic compilation by `setup.py`. ([#482](https://github.com/redballoonsecurity/ofrak/pull/482))
 
 ## [3.2.0](https://github.com/redballoonsecurity/ofrak/compare/ofrak-v3.1.0...ofrak-v3.2.0)
 ### Added

ofrak_core/ofrak/core/entropy/Makefile

@@ -0,0 +1,24 @@
+SHELL := bash
+
+CC = gcc
+CFLAGS = -std=c99 \
+	-pedantic \
+	-Wall \
+	-Wextra \
+	-Werror \
+	-fPIC \
+	-fstack-protector-all \
+	-D_FORTIFY_SOURCE=2 \
+	-shared \
+	-nostdlib \
+	-O3
+LDLIBS = -lm  # Link the math library
+
+# Use this .so.1 extension because otherwise the dependency injector will
+# erroneously try to import entropy.so, which will fail. 
+entropy.so.1: entropy.c
+	$(CC) \
+		$(CFLAGS) \
+		$(filter %.c, $^) \
+		$(LDLIBS) \
+		-o $@

ofrak_core/ofrak/core/entropy/entropy.c

@@ -1,30 +1,27 @@
 #include <stddef.h>   // size_t, NULL
 #include <inttypes.h> // uint8_t, uint32_t
 #include <math.h>     // floor, log2
-// Required to prevent exception with Python >= 3.10
-#define PY_SSIZE_T_CLEAN
-#include <Python.h>
 
 #define HISTOGRAM_SIZE 256
 #define MAX_BRIGHTNESS_FLOAT 255.0
 #define LOGGING_CHUNKS 10
 
-/***
- * Use a Python callback to log the current percent completion of the calculation
- */
-void log_percent(int percent, void* py_callback){
-    PyObject *args = Py_BuildValue("(i)", percent);
-    PyObject *result = PyEval_CallObject(py_callback, args);
-    Py_XDECREF(result);
-    Py_DECREF(args);
-}
+#ifdef _MSC_VER
+    #define EXPORT __declspec(dllexport)
+#else
+    #define EXPORT
+#endif
+
+#ifdef __cplusplus
+extern "C" {
+#endif
 
 /***
  * Calculate the Shannon entropy of a distribution of size `window_size` sampled from a sliding
  * window over `data`. The results of each calculation are stored in `result`.
  */
-int entropy(uint8_t *data, size_t data_len, uint8_t *result, size_t window_size,
-            void* py_log_callback)
+EXPORT int entropy(uint8_t *data, size_t data_len, uint8_t *result, size_t window_size,
+            void (*log_percent)(uint8_t))
 {
     if (data == NULL || result == NULL || window_size > data_len || data_len == 0 ||
         window_size == 0) {
@@ -105,70 +102,12 @@ int entropy(uint8_t *data, size_t data_len, uint8_t *result, size_t window_size,
             }
         }
 
-        log_percent((i * 100) / data_len, py_log_callback);
+        log_percent((i * 100) / data_len);
     }
 
     return 0;
 }
 
-
-PyObject* entropy_wrapper(PyObject* _, PyObject* args){
-    Py_buffer data_buffer;
-    size_t window_size;
-    PyObject* py_log_percent;
-
-    if (!PyArg_ParseTuple(args, "y*nO", &data_buffer, &window_size, &py_log_percent)){
-        PyErr_SetString(PyExc_RuntimeError, "Failed to parse arguments to entropy_wrapper!");
-        return NULL;
-    }
-
-    if (data_buffer.len <= window_size){
-        PyBuffer_Release(&data_buffer);
-         // return b""
-         // we just need a definitely non-NULL pointer to pass to Py_BuildValue
-         // &window_size works fine (no data is read from it)
-        return Py_BuildValue("y#", &window_size, 0);
-    }
-
-    uint8_t *data = data_buffer.buf;
-    size_t result_size = data_buffer.len - window_size;
-    uint8_t *result = (uint8_t*) calloc(result_size, sizeof(uint8_t));
-
-    // Actual entropy calculation
-    entropy(data, data_buffer.len, result, window_size, py_log_percent);
-
-    PyObject* result_object = Py_BuildValue("y#", result, result_size);
-
-    // Clean up memory
-    PyBuffer_Release(&data_buffer);
-    free(result);
-
-    return result_object;
-}
-
-
-// Functions defined in this module
-static PyMethodDef methods[] = {
-    {
-        "entropy_c",
-        entropy_wrapper,
-        METH_VARARGS,
-        "Calculate the Shannon entropy of a distribution of size `window_size` sampled from a sliding window over `data`. The results of each calculation are stored in `result`."
-    },
-    {NULL, NULL, 0, NULL}
-};
-
-
-// Module definition
-static struct PyModuleDef entropy_definition = {
-    PyModuleDef_HEAD_INIT,
-    "entropy_c",
-    "A Python module that calculates Shannon entropy",
-    -1,
-    methods,
-};
-
-PyObject* PyInit_entropy_c(void) {
-    Py_Initialize();
-    return PyModule_Create(&entropy_definition);
+#ifdef __cplusplus
 }
+#endif

ofrak_core/ofrak/core/entropy/entropy.py

@@ -1,5 +1,4 @@
 import asyncio
-import ctypes
 import logging
 import math
 from concurrent.futures import ProcessPoolExecutor
@@ -14,10 +13,6 @@
 
 LOGGER = logging.getLogger(__name__)
 
-
-C_LOG_TYPE = ctypes.CFUNCTYPE(None, ctypes.c_uint8)
-
-
 try:
     from ofrak.core.entropy.entropy_c import entropy_c as entropy_func
 except:

ofrak_core/ofrak/core/entropy/entropy_c.py

@@ -0,0 +1,39 @@
+import ctypes
+import os
+from typing import Callable, Optional
+import sysconfig
+
+C_LOG_TYPE = ctypes.CFUNCTYPE(None, ctypes.c_uint8)
+
+ext_suffix = sysconfig.get_config_var("EXT_SUFFIX")
+if not isinstance(ext_suffix, str):
+    raise RuntimeError("Could not find compiled C library, no EXT_SUFFIX sysconfig var")
+
+_lib_entropy = ctypes.cdll.LoadLibrary(
+    os.path.abspath(os.path.join(os.path.dirname(__file__), "entropy_c" + ext_suffix + ".1"))
+)
+C_ENTROPY_FUNC = _lib_entropy.entropy
+
+C_ENTROPY_FUNC.argtypes = (
+    ctypes.c_char_p,
+    ctypes.c_size_t,
+    ctypes.c_char_p,
+    ctypes.c_size_t,
+    C_LOG_TYPE,
+)
+C_ENTROPY_FUNC.restype = ctypes.c_int
+
+
+def entropy_c(
+    data: bytes, window_size: int, log_percent: Optional[Callable[[int], None]] = None
+) -> bytes:
+    if log_percent is None:
+        log_percent = lambda x: None
+
+    if len(data) <= window_size:
+        return b""
+    entropy = ctypes.create_string_buffer(len(data) - window_size)
+    errval = C_ENTROPY_FUNC(data, len(data), entropy, window_size, C_LOG_TYPE(log_percent))
+    if errval != 0:
+        raise ValueError("Bad input to entropy function.")
+    return bytes(entropy.raw)

ofrak_core/setup.py

@@ -1,6 +1,8 @@
+import sys
 import setuptools
 import pkg_resources
 from setuptools.command.egg_info import egg_info
+from setuptools.command.build_ext import build_ext
 
 
 class egg_info_ex(egg_info):
@@ -16,16 +18,43 @@ def run(self):
         egg_info.run(self)
 
 
+class build_ext_1(build_ext):
+    """Changes the output filename of ctypes libraries to have '.1' at the end
+    so they don't interfere with the dependency injection.
+
+    Based on: https://stackoverflow.com/a/34830639
+    """
+
+    def get_export_symbols(self, ext):
+        if isinstance(ext, CTypesExtension):
+            return ext.export_symbols
+        return super().get_export_symbols(ext)
+
+    def get_ext_filename(self, ext_name):
+        default_filename = super().get_ext_filename(ext_name)
+
+        if ext_name in self.ext_map:
+            ext = self.ext_map[ext_name]
+            if isinstance(ext, CTypesExtension):
+                return default_filename + ".1"
+
+        return default_filename
+
+
+class CTypesExtension(setuptools.Extension):
+    pass
+
+
 with open("README.md") as f:
     long_description = f.read()
 
 
-entropy_so = setuptools.Extension(
+entropy_so = CTypesExtension(
     "ofrak.core.entropy.entropy_c",
     sources=["ofrak/core/entropy/entropy.c"],
-    libraries=["m"],  # math library
+    libraries=["m"] if sys.platform != "win32" else None,  # math library
     optional=True,  # If this fails the build, OFRAK will fall back to Python implementation
-    extra_compile_args=["-O3"],
+    extra_compile_args=["-O3"] if sys.platform != "win32" else ["/O2"],
 )
 
 
@@ -84,7 +113,7 @@ def read_requirements(requirements_path):
     python_requires=">=3.7",
     license="Proprietary",
     license_files=["LICENSE"],
-    cmdclass={"egg_info": egg_info_ex},
+    cmdclass={"egg_info": egg_info_ex, "build_ext": build_ext_1},
     entry_points={
         "ofrak.packages": ["ofrak_pkg = ofrak"],
         "console_scripts": ["ofrak = ofrak.__main__:main"],