new: implemented artifacts extraction and saving (closes #24)

.gitignore

@@ -1,7 +1,7 @@
 /trace.json
 /linux_test_executable
 /malicious.*
-
+/artifacts
 # Testing code
 notebooks/
 

README.md

@@ -79,6 +79,12 @@ Create a trace file for a given loader with:
 dyana trace --loader automodel ... --output trace.json
 ```
 
+To save artifacts from the container, you can pass the `--save` flag:
+
+```bash
+dyana trace --loader pip --package botocore --save /usr/local/bin/jp.py --save-to ./artifacts
+```
+
 It is possible to override the default events that Dyana will trace by passing a [custom policy](https://aquasecurity.github.io/tracee/v0.14/docs/policies/) to the tracer with:
 
 ```bash

dyana/cli.py

@@ -100,6 +100,8 @@ def trace(
     loader: str = typer.Option(help="Loader to use.", default="automodel"),
     platform: str | None = typer.Option(help="Platform to use.", default=None),
     output: pathlib.Path = typer.Option(help="Path to the output file.", default="trace.json"),
+    save: list[str] = typer.Option(help="List of file artifacts to save.", default=[]),
+    save_to: pathlib.Path = typer.Option(help="Path to the directory to save the artifacts to.", default="./artifacts"),
     policy: pathlib.Path | None = typer.Option(
         help="Path to a policy or directory with custom tracee policies.", default=None
     ),
@@ -118,7 +120,9 @@ def trace(
         if policy and not policy.exists():
             raise typer.BadParameter(f"policy file or directory not found: {policy}")
 
-        the_loader = Loader(name=loader, timeout=timeout, platform=platform, args=ctx.args, verbose=verbose)
+        the_loader = Loader(
+            name=loader, timeout=timeout, platform=platform, args=ctx.args, verbose=verbose, save=save, save_to=save_to
+        )
         the_tracer = Tracer(the_loader, policy=policy)
 
         trace = the_tracer.run_trace(allow_network, not no_gpu, allow_volume_write)

dyana/docker.py

@@ -115,6 +115,7 @@ def run_detached(
     image: str,
     command: list[str],
     volumes: dict[str, str],
+    environment: dict[str, str] | None = None,
     allow_network: bool = False,
     allow_gpus: bool = True,
     allow_volume_write: bool = False,
@@ -125,7 +126,10 @@ def run_detached(
     network_mode = "bridge" if allow_network else "none"
 
     # by default volumes are read-only
-    mounts = {host: {"bind": guest, "mode": "rw" if allow_volume_write else "ro"} for host, guest in volumes.items()}
+    mounts = {
+        host: {"bind": guest, "mode": "rw" if guest == "/artifacts" or allow_volume_write else "ro"}
+        for host, guest in volumes.items()
+    }
 
     # this allows us to log dns requests even if the container is in network mode "none"
     dns = ["127.0.0.1"] if not allow_network else None
@@ -138,6 +142,7 @@ def run_detached(
         image,
         command=command,
         volumes=mounts,
+        environment=environment,
         network_mode=network_mode,
         dns=dns,
         # automatically remove the container after it exits

dyana/loaders/base/dyana.py

@@ -1,3 +1,6 @@
+import atexit
+import os
+import pathlib
 import resource
 import shutil
 import sys
@@ -6,6 +9,22 @@
 from io import StringIO
 
 
+def save_artifacts() -> None:
+    artifacts = os.environ.get("DYANA_SAVE", "").split(",")
+    if artifacts:
+        for artifact in artifacts:
+            try:
+                if os.path.isdir(artifact):
+                    shutil.copytree(artifact, f"/artifacts/{artifact}")
+                elif os.path.isfile(artifact):
+                    shutil.copy(artifact, "/artifacts")
+            except Exception:
+                pass
+
+
+atexit.register(save_artifacts)
+
+
 class Profiler:
     def __init__(self, gpu: bool = False):
         self._errors: dict[str, str] = {}

dyana/loaders/loader.py

@@ -47,6 +47,8 @@ def __init__(
         build: bool = True,
         platform: str | None = None,
         args: list[str] | None = None,
+        save: list[str] = [],
+        save_to: pathlib.Path = pathlib.Path("./artifacts"),
         verbose: bool = False,
     ):
         # make sure that name does not include a path traversal
@@ -66,6 +68,8 @@ def __init__(
         self.settings: LoaderSettings | None = None
         self.build_args: dict[str, str] | None = None
         self.args: list[ParsedArgument] | None = None
+        self.save: list[str] = save
+        self.save_to: pathlib.Path = save_to.resolve().absolute()
 
         if os.path.exists(self.settings_path):
             with open(self.settings_path) as f:
@@ -180,8 +184,23 @@ def run(self, allow_network: bool = False, allow_gpus: bool = True, allow_volume
 
         try:
             self.output = ""
+            environment = {}
+            if self.save:
+                environment["DYANA_SAVE"] = ",".join(self.save)
+                volumes[str(self.save_to)] = "/artifacts"
+                if not os.path.exists(self.save_to):
+                    os.makedirs(self.save_to)
+
+                print(f":popcorn: [bold]loader[/]: saving artifacts to [dim]{self.save_to}[/]")
+
             self.container = docker.run_detached(
-                self.image, arguments, volumes, allow_network, allow_gpus, allow_volume_write
+                self.image,
+                arguments,
+                volumes,
+                environment=environment,
+                allow_network=allow_network,
+                allow_gpus=allow_gpus,
+                allow_volume_write=allow_volume_write,
             )
             self.container_id = self.container.id
             self.reader_thread = threading.Thread(target=self._reader_thread)