From c46abbcec306b9845ab32b7282090c9f01629df4 Mon Sep 17 00:00:00 2001
From: bmc-msft <41130664+bmc-msft@users.noreply.github.com>
Date: Wed, 2 Jun 2021 13:29:20 -0400
Subject: [PATCH] stabilize `onefuzz jobs containers download` (#953)

---
 src/cli/onefuzz/api.py   | 28 ++++++++++++++++++++++++++++
 src/cli/onefuzz/debug.py | 22 +---------------------
 2 files changed, 29 insertions(+), 21 deletions(-)

diff --git a/src/cli/onefuzz/api.py b/src/cli/onefuzz/api.py
index 8ecf47fa64..5342fd39fd 100644
--- a/src/cli/onefuzz/api.py
+++ b/src/cli/onefuzz/api.py
@@ -30,6 +30,7 @@
 from six.moves import input  # workaround for static analysis
 
 from .__version__ import __version__
+from .azcopy import azcopy_sync
 from .backend import Backend, BackendConfig, ContainerWrapper, wait
 from .ssh import build_ssh_command, ssh_connect, temp_file
 
@@ -970,6 +971,33 @@ def list(
             results[container] = self.onefuzz.containers.files.list(container).files
         return results
 
+    def download(
+        self, job_id: UUID_EXPANSION, *, output: Optional[primitives.Directory] = None
+    ) -> None:
+        to_download = {}
+        tasks = self.onefuzz.tasks.list(job_id=job_id, state=None)
+        if not tasks:
+            raise Exception("no tasks with job_id:%s" % job_id)
+
+        for task in tasks:
+            for container in task.config.containers:
+                info = self.onefuzz.containers.get(container.name)
+                name = os.path.join(container.type.name, container.name)
+                to_download[name] = info.sas_url
+
+        if output is None:
+            output = primitives.Directory(os.getcwd())
+
+        for name in to_download:
+            outdir = os.path.join(output, name)
+            if not os.path.exists(outdir):
+                os.makedirs(outdir)
+            self.logger.info("downloading: %s", name)
+            # 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(to_download[name], outdir)
+
     def delete(
         self,
         job_id: UUID_EXPANSION,
diff --git a/src/cli/onefuzz/debug.py b/src/cli/onefuzz/debug.py
index 77869eb195..239d4e06fa 100644
--- a/src/cli/onefuzz/debug.py
+++ b/src/cli/onefuzz/debug.py
@@ -22,7 +22,6 @@
 
 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
@@ -337,26 +336,7 @@ def libfuzzer_execs_sec(
     def download_files(self, job_id: UUID_EXPANSION, output: Directory) -> None:
         """Download the containers by container type for each task in the specified job"""
 
-        to_download = {}
-        tasks = self.onefuzz.tasks.list(job_id=job_id, state=None)
-        if not tasks:
-            raise Exception("no tasks with job_id:%s" % job_id)
-
-        for task in tasks:
-            for container in task.config.containers:
-                info = self.onefuzz.containers.get(container.name)
-                name = os.path.join(container.type.name, container.name)
-                to_download[name] = info.sas_url
-
-        for name in to_download:
-            outdir = os.path.join(output, name)
-            if not os.path.exists(outdir):
-                os.makedirs(outdir)
-            self.logger.info("downloading: %s", name)
-            # 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(to_download[name], outdir)
+        self.onefuzz.jobs.containers.download(job_id, output=output)
 
 
 class DebugLog(Command):