feat: add safetensors loader for structural integrity analysis

README.md

@@ -28,7 +28,7 @@
     </a>
 </h4>
 
-Dyana is a sandbox environment using Docker and [Tracee](https://github.com/aquasecurity/tracee) for loading, running and profiling a wide range of files, including machine learning models, ELF executables, Pickle serialized files, Javascripts [and more](https://docs.dreadnode.io/open-source/dyana/topics/loaders). It provides detailed insights into GPU memory usage, filesystem interactions, network requests, and security related events.
+Dyana is a sandbox environment using Docker and [Tracee](https://github.com/aquasecurity/tracee) for loading, running and profiling a wide range of files, including machine learning models, SafeTensors model files, ELF executables, Pickle serialized files, Javascripts [and more](https://docs.dreadnode.io/open-source/dyana/topics/loaders). It provides detailed insights into GPU memory usage, filesystem interactions, network requests, and security related events.
 
 ## Installation
 

dyana/cli.py

@@ -20,6 +20,7 @@
 from dyana.view import (
     view_disk_events,
     view_disk_usage,
+    view_extra,
     view_gpus,
     view_header,
     view_imports,
@@ -139,7 +140,9 @@ def trace(
     except Exception as e:
         serr = str(e)
         if "could not select device driver" in serr and "capabilities: [[gpu]]" in serr:
-            rich_print(":cross_mark: [bold][red]error:[/] [red]GPUs are not available on this system, run with --no-gpu.[/]")
+            rich_print(
+                ":cross_mark: [bold][red]error:[/] [red]GPUs are not available on this system, run with --no-gpu.[/]"
+            )
         else:
             rich_print(f":cross_mark: [bold][red]error:[/] [red]{e}[/]")
 
@@ -187,3 +190,4 @@ def summary(trace_path: pathlib.Path = typer.Option(help="Path to the trace file
         view_legacy_extra(trace["run"])
     else:
         view_imports(trace["run"]["stages"])
+        view_extra(trace["run"])

dyana/cli_test.py

@@ -177,9 +177,7 @@ def test_trace_runs_and_saves(self, _mock_run: t.Any, tmp_path: t.Any) -> None:
     @patch("dyana.cli.Tracer.__init__", _noop_tracer_init)
     @patch(
         "dyana.cli.Tracer.run_trace",
-        side_effect=RuntimeError(
-            "could not select device driver '' with capabilities: [[gpu]]"
-        ),
+        side_effect=RuntimeError("could not select device driver '' with capabilities: [[gpu]]"),
     )
     @patch("dyana.cli.Loader.__init__", _noop_loader_init)
     def test_trace_gpu_error(self, _mock_run: t.Any) -> None:

dyana/loaders/base/dyana_test.py

@@ -18,62 +18,48 @@ def setup_method(self) -> None:
 
     def test_singleton(self) -> None:
         with patch("dyana.loaders.base.dyana.Stage.create") as mock_create:
-            mock_create.return_value = Stage(
-                name="start", timestamp=0, ram=0, disk=0, network={}, imports={}
-            )
+            mock_create.return_value = Stage(name="start", timestamp=0, ram=0, disk=0, network={}, imports={})
             p = Profiler()
             assert Profiler.instance is p
 
     def test_on_stage(self) -> None:
         with patch("dyana.loaders.base.dyana.Stage.create") as mock_create:
-            mock_create.return_value = Stage(
-                name="test", timestamp=0, ram=0, disk=0, network={}, imports={}
-            )
+            mock_create.return_value = Stage(name="test", timestamp=0, ram=0, disk=0, network={}, imports={})
             p = Profiler()
             p.on_stage("after_load")
             assert len(p._stages) == 2
 
     def test_track_error(self) -> None:
         with patch("dyana.loaders.base.dyana.Stage.create") as mock_create:
-            mock_create.return_value = Stage(
-                name="start", timestamp=0, ram=0, disk=0, network={}, imports={}
-            )
+            mock_create.return_value = Stage(name="start", timestamp=0, ram=0, disk=0, network={}, imports={})
             p = Profiler()
             p.track_error("loader", "something broke")
             assert p._errors == {"loader": "something broke"}
 
     def test_track_warning(self) -> None:
         with patch("dyana.loaders.base.dyana.Stage.create") as mock_create:
-            mock_create.return_value = Stage(
-                name="start", timestamp=0, ram=0, disk=0, network={}, imports={}
-            )
+            mock_create.return_value = Stage(name="start", timestamp=0, ram=0, disk=0, network={}, imports={})
             p = Profiler()
             p.track_warning("pip", "could not import")
             assert p._warnings == {"pip": "could not import"}
 
     def test_track_extra(self) -> None:
         with patch("dyana.loaders.base.dyana.Stage.create") as mock_create:
-            mock_create.return_value = Stage(
-                name="start", timestamp=0, ram=0, disk=0, network={}, imports={}
-            )
+            mock_create.return_value = Stage(name="start", timestamp=0, ram=0, disk=0, network={}, imports={})
             p = Profiler()
             p.track_extra("imports", {"os": "/usr/lib"})
             assert p._extra == {"imports": {"os": "/usr/lib"}}
 
     def test_track(self) -> None:
         with patch("dyana.loaders.base.dyana.Stage.create") as mock_create:
-            mock_create.return_value = Stage(
-                name="start", timestamp=0, ram=0, disk=0, network={}, imports={}
-            )
+            mock_create.return_value = Stage(name="start", timestamp=0, ram=0, disk=0, network={}, imports={})
             p = Profiler()
             p.track("custom_key", "custom_value")
             assert p._additionals == {"custom_key": "custom_value"}
 
     def test_as_dict(self) -> None:
         with patch("dyana.loaders.base.dyana.Stage.create") as mock_create:
-            mock_create.return_value = Stage(
-                name="start", timestamp=0, ram=0, disk=0, network={}, imports={}
-            )
+            mock_create.return_value = Stage(name="start", timestamp=0, ram=0, disk=0, network={}, imports={})
             p = Profiler()
             p.track_error("err", "msg")
             result = p.as_dict()
@@ -85,9 +71,7 @@ def test_as_dict(self) -> None:
 
     def test_flush(self, capsys: t.Any) -> None:
         with patch("dyana.loaders.base.dyana.Stage.create") as mock_create:
-            mock_create.return_value = Stage(
-                name="start", timestamp=0, ram=0, disk=0, network={}, imports={}
-            )
+            mock_create.return_value = Stage(name="start", timestamp=0, ram=0, disk=0, network={}, imports={})
             Profiler()
             Profiler.flush()
             captured = capsys.readouterr()
@@ -151,7 +135,9 @@ def test_with_prev_imports(self) -> None:
             patch("dyana.loaders.base.dyana.get_peak_rss", return_value=1024),
             patch("dyana.loaders.base.dyana.get_disk_usage", return_value=2048),
             patch("dyana.loaders.base.dyana.get_network_stats", return_value={}),
-            patch("dyana.loaders.base.dyana.get_current_imports", return_value={"os": "/a", "sys": "/b", "new_mod": "/c"}),
+            patch(
+                "dyana.loaders.base.dyana.get_current_imports", return_value={"os": "/a", "sys": "/b", "new_mod": "/c"}
+            ),
         ):
             stage = Stage.create("test", prev_imports={"os": "/a", "sys": "/b"})
             assert "new_mod" in stage.imports

dyana/loaders/loader.py

@@ -184,7 +184,9 @@ def run(self, allow_network: bool = False, allow_gpus: bool = True, allow_volume
             rich_print(":popcorn: [bold]loader[/]: [yellow]required bridged network access[/]")
 
         elif allow_network:
-            rich_print(":popcorn: [bold]loader[/]: [yellow]warning: allowing bridged network access to the container[/]")
+            rich_print(
+                ":popcorn: [bold]loader[/]: [yellow]warning: allowing bridged network access to the container[/]"
+            )
 
         if allow_volume_write:
             rich_print(":popcorn: [bold]loader[/]: [yellow]warning: allowing volume write to the container[/]")

dyana/loaders/safetensors/.gitignore

@@ -0,0 +1,3 @@
+dyana.py
+dyana-requirements.txt
+dyana-requirements-gpu.txt

dyana/loaders/safetensors/Dockerfile

@@ -0,0 +1,15 @@
+FROM python:3.12-slim
+
+WORKDIR /app
+
+RUN apt-get update && apt-get install -y build-essential
+COPY dyana.py .
+COPY dyana-requirements.txt .
+RUN pip install --no-cache-dir --root-user-action=ignore -r dyana-requirements.txt
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir --root-user-action=ignore -r requirements.txt
+
+COPY main.py .
+
+ENTRYPOINT ["python3", "-W", "ignore", "main.py"]