diff --git a/src/cli/onefuzz/api.py b/src/cli/onefuzz/api.py index 0038855b6..5ccc53c67 100644 --- a/src/cli/onefuzz/api.py +++ b/src/cli/onefuzz/api.py @@ -460,6 +460,44 @@ class Containers(Endpoint): self.logger.debug("list containers") return self._req_model_list("GET", responses.ContainerInfoBase) + def download_job( + self, job_id: UUID_EXPANSION, *, output: Optional[primitives.Directory] = None + ) -> None: + tasks = self.onefuzz.tasks.list(job_id=job_id, state=None) + if not tasks: + raise Exception("no tasks with job_id:%s" % job_id) + + self._download_tasks(tasks, output) + + def download_task( + self, task_id: UUID_EXPANSION, *, output: Optional[primitives.Directory] = None + ) -> None: + self._download_tasks([self.onefuzz.tasks.get(task_id=task_id)], output) + + def _download_tasks( + self, tasks: List[models.Task], output: Optional[primitives.Directory] + ) -> None: + + to_download: Dict[str, str] = {} + 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) + class Repro(Endpoint): """Interact with Reproduction VMs""" @@ -1005,33 +1043,6 @@ class JobContainers(Endpoint): 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 cd446eaf2..29ade1008 100644 --- a/src/cli/onefuzz/debug.py +++ b/src/cli/onefuzz/debug.py @@ -174,7 +174,7 @@ class DebugScaleset(Command): class DebugTask(Command): - """Debug a specific job""" + """Debug a specific task""" def list_nodes(self, task_id: UUID_EXPANSION) -> Optional[List[NodeAssignment]]: task = self.onefuzz.tasks.get(task_id) @@ -255,6 +255,11 @@ class DebugTask(Command): query, timespan, limit ) + def download_files(self, task_id: UUID_EXPANSION, output: Directory) -> None: + """Download the containers by container type for the specified task""" + + self.onefuzz.containers.download_task(task_id, output=output) + class DebugJobTask(Command): """Debug a task for a specific job""" @@ -339,7 +344,7 @@ 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""" - self.onefuzz.jobs.containers.download(job_id, output=output) + self.onefuzz.containers.download_job(job_id, output=output) def rerun( self,