Added semgrep and some fixes

luxonis · Feb 24, 2025 · 7b430a5 · 7b430a5
1 parent e37a315
commit 7b430a5
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 6 deletions.
diff --git a/modelconverter/packages/hailo/exporter.py b/modelconverter/packages/hailo/exporter.py
@@ -21,8 +21,8 @@
 
 @contextmanager
 def _replace_module(original, substitute):
-    original_module = importlib.import_module(original)
-    substitute_module = importlib.import_module(substitute)
+    original_module = importlib.import_module(original)  # nosemgrep
+    substitute_module = importlib.import_module(substitute)  # nosemgrep
 
     sys.modules[original] = substitute_module
     try:

diff --git a/modelconverter/packages/multistage_exporter.py b/modelconverter/packages/multistage_exporter.py
@@ -110,10 +110,21 @@ def _produce_calibration_data(self, exporter: Exporter) -> None:
                         for out_name in output_dirs
                     }
 
-                    # TODO: safe exec
                     local_scope = {}
+                    safe_globals = {"__builtins__": {}}
+
+                    try:
+                        exec(  # nosemgrep
+                            script, safe_globals, local_scope
+                        )
+                    except Exception as e:
+                        raise RuntimeError(f"Error executing script: {e}")
+
+                    if "run_script" not in local_scope:
+                        raise RuntimeError(
+                            "Error: `run_script` function not found in script."
+                        )
 
-                    exec(script, globals(), local_scope)
                     run_script = local_scope["run_script"]
                     arr = run_script(outputs)
                     np.save(dest / f"{i}.npy", arr)

diff --git a/modelconverter/utils/nn_archive.py b/modelconverter/utils/nn_archive.py
@@ -47,8 +47,21 @@ def process_nn_archive(
     elif tarfile.is_tarfile(path):
         if untar_path.suffix == ".tar":
             untar_path = MISC_DIR / untar_path.stem
-        with tarfile.open(path) as tar:
-            tar.extractall(untar_path)
+
+        def safe_members(tar):
+            """Filter members to prevent path traversal attacks."""
+            safe_files = []
+            for member in tar.getmembers():
+                # Normalize path and ensure it's within the extraction folder
+                if not member.name.startswith("/") and ".." not in member.name:
+                    safe_files.append(member)
+                else:
+                    logger.warning(f"Skipping unsafe file: {member.name}")
+            return safe_files
+
+        tf = tarfile.open(path, mode="r")
+        tf.extractall(untar_path, members=safe_members(tf))
+
     else:
         raise RuntimeError(f"Unknown NN Archive path: `{path}`")