diff --git a/src/cli/onefuzz/api.py b/src/cli/onefuzz/api.py index 02832f24e..1d4b74e54 100644 --- a/src/cli/onefuzz/api.py +++ b/src/cli/onefuzz/api.py @@ -215,7 +215,9 @@ class Files(Endpoint): client = self._get_client(container) client.upload_file(file_path, blob_name) - def upload_dir(self, container: primitives.Container, dir_path: str) -> None: + def upload_dir( + self, container: primitives.Container, dir_path: primitives.Directory + ) -> None: """ uploads a directory to a container """ self.logger.debug("uploading directory to container %s:%s", container, dir_path) @@ -223,6 +225,18 @@ class Files(Endpoint): client = self._get_client(container) client.upload_dir(dir_path) + def download_dir( + self, container: primitives.Container, dir_path: primitives.Directory + ) -> None: + """ downloads a container to a directory """ + + self.logger.debug( + "downloading container to directory %s:%s", container, dir_path + ) + + client = self._get_client(container) + client.download_dir(dir_path) + class Versions(Endpoint): """ Onefuzz Instance """ diff --git a/src/cli/onefuzz/azcopy.py b/src/cli/onefuzz/azcopy.py new file mode 100644 index 000000000..014eb956a --- /dev/null +++ b/src/cli/onefuzz/azcopy.py @@ -0,0 +1,16 @@ +import os +import shutil +import subprocess # nosec + + +def azcopy_sync(src: str, dst: str) -> None: + """ Expose azcopy for uploading/downloading files """ + + azcopy = os.environ.get("AZCOPY") or shutil.which("azcopy") + if not azcopy: + raise Exception( + "unable to find 'azcopy' in path or AZCOPY environment variable" + ) + + # security note: callers need to understand the src/dst for this. + subprocess.check_output([azcopy, "sync", src, dst]) # nosec diff --git a/src/cli/onefuzz/backend.py b/src/cli/onefuzz/backend.py index 1e8ccaf27..a8d42b7e2 100644 --- a/src/cli/onefuzz/backend.py +++ b/src/cli/onefuzz/backend.py @@ -38,6 +38,8 @@ from tenacity.retry import retry_if_exception_type from tenacity.stop import stop_after_attempt from tenacity.wait import wait_random +from .azcopy import azcopy_sync + _ACCESSTOKENCACHE_UMASK = 0o077 ONEFUZZ_BASE_PATH = os.path.join("~", ".cache", "onefuzz") @@ -294,6 +296,7 @@ def before_sleep( class ContainerWrapper: def __init__(self, container_url: str) -> None: self.client = ContainerClient.from_container_url(container_url) + self.container_url = container_url @retry( stop=stop_after_attempt(10), @@ -320,6 +323,12 @@ class ContainerWrapper: blob_name = os.path.relpath(path, start=dir_path) self.upload_file(path, blob_name) + def download_dir(self, dir_path: str) -> None: + # security note: the src for azcopy comes from the server which is + # trusted in this context, while the destination is provided by the + # user + azcopy_sync(self.container_url, dir_path) + @retry( stop=stop_after_attempt(10), wait=wait_random(min=1, max=3), diff --git a/src/cli/onefuzz/debug.py b/src/cli/onefuzz/debug.py index 568d5f92f..f7f28dd83 100644 --- a/src/cli/onefuzz/debug.py +++ b/src/cli/onefuzz/debug.py @@ -6,8 +6,6 @@ import json import logging import os -import shutil -import subprocess # nosec import tempfile import time from typing import Any, Dict, List, Optional, Tuple, Union @@ -24,6 +22,7 @@ from onefuzztypes.primitives import Container, Directory from onefuzz.api import UUID_EXPANSION, Command, Onefuzz +from .azcopy import azcopy_sync from .backend import wait from .rdp import rdp_connect from .ssh import ssh_connect @@ -338,12 +337,6 @@ class DebugJob(Command): def download_files(self, job_id: UUID_EXPANSION, output: Directory) -> None: """ Download the containers by container type for each task in the specified job """ - azcopy = os.environ.get("AZCOPY") or shutil.which("azcopy") - if not azcopy: - raise Exception( - "unable to find 'azcopy' in path or AZCOPY environment variable" - ) - to_download = {} tasks = self.onefuzz.tasks.list(job_id=job_id, state=None) if not tasks: @@ -363,9 +356,7 @@ class DebugJob(Command): # security note: the src for azcopy comes from the server which is # trusted in this context, while the destination is provided by the # user - subprocess.check_output( # nosec - [azcopy, "sync", to_download[name], outdir] - ) + azcopy_sync(to_download[name], outdir) class DebugLog(Command): diff --git a/src/cli/onefuzz/templates/__init__.py b/src/cli/onefuzz/templates/__init__.py index a16007f21..976ce4777 100644 --- a/src/cli/onefuzz/templates/__init__.py +++ b/src/cli/onefuzz/templates/__init__.py @@ -200,7 +200,7 @@ class JobHelper: self.logger.info("uploading inputs from zip: `%s`" % path) self.onefuzz.containers.files.upload_dir( - self.containers[ContainerType.inputs], tmp_dir + self.containers[ContainerType.inputs], Directory(tmp_dir) ) @classmethod