From dbc79ed5d4fa2d6b7d17ed9b1f835f7d563b4f20 Mon Sep 17 00:00:00 2001 From: evilsocket Date: Sat, 18 Jan 2025 17:35:23 +0100 Subject: [PATCH] new: implemented artifacts extraction and saving (closes #24) --- .gitignore | 2 +- README.md | 6 ++++++ dyana/cli.py | 6 +++++- dyana/docker.py | 7 ++++++- dyana/loaders/base/dyana.py | 19 +++++++++++++++++++ dyana/loaders/loader.py | 21 ++++++++++++++++++++- 6 files changed, 57 insertions(+), 4 deletions(-) diff --git a/.gitignore b/.gitignore index f32b232..e565844 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,7 @@ /trace.json /linux_test_executable /malicious.* - +/artifacts # Testing code notebooks/ diff --git a/README.md b/README.md index dca1b9c..7e48ebe 100644 --- a/README.md +++ b/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 diff --git a/dyana/cli.py b/dyana/cli.py index a687c3f..6800a01 100644 --- a/dyana/cli.py +++ b/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) diff --git a/dyana/docker.py b/dyana/docker.py index 7ee4714..41aa20e 100644 --- a/dyana/docker.py +++ b/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 diff --git a/dyana/loaders/base/dyana.py b/dyana/loaders/base/dyana.py index a6b4752..702818e 100644 --- a/dyana/loaders/base/dyana.py +++ b/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] = {} diff --git a/dyana/loaders/loader.py b/dyana/loaders/loader.py index 9dd8e93..714d74f 100644 --- a/dyana/loaders/loader.py +++ b/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)