Add endpoint to download the agent binaries (#2600)

* Add endpoint to download the agent binaries

* build fix

* fix types

* adding test

* format

* remove comments

* address pr comment

* build fix

src/ApiService/ApiService/Functions/Tool.cs

@@ -0,0 +1,31 @@
+using System.Net;
+using Microsoft.Azure.Functions.Worker;
+using Microsoft.Azure.Functions.Worker.Http;
+
+namespace Microsoft.OneFuzz.Service.Functions;
+
+public class Tools {
+    private readonly IOnefuzzContext _context;
+    private readonly IEndpointAuthorization _auth;
+
+    public Tools(IEndpointAuthorization auth, IOnefuzzContext context) {
+        _context = context;
+        _auth = auth;
+    }
+
+    public async Async.Task<HttpResponseData> GetResponse(HttpRequestData req) {
+        //Note: streaming response are not currently supported by in isolated functions
+        // https://github.com/Azure/azure-functions-dotnet-worker/issues/958
+        var response = req.CreateResponse(HttpStatusCode.OK);
+        var downloadResult = await _context.Containers.DownloadAsZip(WellKnownContainers.Tools, StorageType.Config, response.Body);
+        if (!downloadResult.IsOk) {
+            return await _context.RequestHandling.NotOk(req, downloadResult.ErrorV, "download tools");
+        }
+        return response;
+    }
+
+
+    [Function("Tools")]
+    public Async.Task<HttpResponseData> Run([HttpTrigger(AuthorizationLevel.Anonymous, "GET")] HttpRequestData req)
+        => _auth.CallIfUser(req, GetResponse);
+}

src/ApiService/ApiService/OneFuzzTypes/Enums.cs

@@ -29,6 +29,7 @@ public enum ErrorCode {
     PROXY_FAILED = 472,
     INVALID_CONFIGURATION = 473,
     UNABLE_TO_CREATE_CONTAINER = 474,
+    UNABLE_TO_DOWNLOAD_FILE = 475,
 }
 
 public enum VmState {

src/ApiService/ApiService/onefuzzlib/Containers.cs

@@ -1,8 +1,11 @@
-using System.Threading;
+using System.IO;
+using System.IO.Compression;
+using System.Threading;
 using System.Threading.Tasks;
 using Azure;
 using Azure.Storage.Blobs;
 using Azure.Storage.Blobs.Models;
+using Azure.Storage.Blobs.Specialized;
 using Azure.Storage.Sas;
 
 namespace Microsoft.OneFuzz.Service;
@@ -31,6 +34,7 @@ public interface IContainers {
     public Async.Task<Dictionary<Container, IDictionary<string, string>>> GetContainers(StorageType corpus);
 
     public string AuthDownloadUrl(Container container, string filename);
+    public Async.Task<OneFuzzResultVoid> DownloadAsZip(Container container, StorageType storageType, Stream stream, string? prefix = null);
 }
 
 public class Containers : IContainers {
@@ -234,4 +238,21 @@ public class Containers : IContainers {
 
         return $"{instance}/api/download?{queryString}";
     }
+
+    public async Async.Task<OneFuzzResultVoid> DownloadAsZip(Container container, StorageType storageType, Stream stream, string? prefix = null) {
+        var client = await FindContainer(container, storageType) ?? throw new Exception($"unable to find container: {container} - {storageType}");
+        var blobs = client.GetBlobs(prefix: prefix);
+
+        using var archive = new ZipArchive(stream, ZipArchiveMode.Create, true);
+        await foreach (var b in blobs.ToAsyncEnumerable()) {
+            var entry = archive.CreateEntry(b.Name);
+            await using var entryStream = entry.Open();
+            var blobClient = client.GetBlockBlobClient(b.Name);
+            var downloadResult = await blobClient.DownloadToAsync(entryStream);
+            if (downloadResult.IsError) {
+                return OneFuzzResultVoid.Error(ErrorCode.UNABLE_TO_DOWNLOAD_FILE, $"Error while downloading blob {b.Name}");
+            }
+        }
+        return OneFuzzResultVoid.Ok;
+    }
 }

src/ApiService/IntegrationTests/ToolsTests.cs

@@ -0,0 +1,66 @@
+using System;
+using System.IO;
+using System.IO.Compression;
+using System.Linq;
+using System.Net;
+using IntegrationTests.Fakes;
+using Microsoft.OneFuzz.Service;
+using Microsoft.OneFuzz.Service.Functions;
+using Xunit;
+using Xunit.Abstractions;
+using Async = System.Threading.Tasks;
+
+namespace IntegrationTests;
+
+[Trait("Category", "Live")]
+public class AzureStorageToolsTest : ToolsTestBase {
+    public AzureStorageToolsTest(ITestOutputHelper output)
+        : base(output, Integration.AzureStorage.FromEnvironment()) { }
+}
+
+public class AzuriteToolsTest : ToolsTestBase {
+    public AzuriteToolsTest(ITestOutputHelper output)
+        : base(output, new Integration.AzuriteStorage()) { }
+}
+
+public abstract class ToolsTestBase : FunctionTestBase {
+    private readonly IStorage _storage;
+
+    public ToolsTestBase(ITestOutputHelper output, IStorage storage)
+        : base(output, storage) {
+        _storage = storage;
+    }
+
+    [Fact]
+    public async Async.Task CanDownload() {
+
+        const int NUMBER_OF_FILES = 20;
+
+        var toolsContainerClient = GetContainerClient(WellKnownContainers.Tools);
+        _ = await toolsContainerClient.CreateIfNotExistsAsync();
+
+        // generate random content 
+        var files = Enumerable.Range(0, NUMBER_OF_FILES).Select((x, i) => (path: i, content: Guid.NewGuid())).ToList();
+
+        // upload each files
+        foreach (var (path, content) in files) {
+            var r = await toolsContainerClient.UploadBlobAsync(path.ToString(), BinaryData.FromString(content.ToString()));
+            Assert.False(r.GetRawResponse().IsError);
+        }
+        var auth = new TestEndpointAuthorization(RequestType.User, Logger, Context);
+        var func = new Tools(auth, Context);
+        var result = await func.Run(TestHttpRequestData.FromJson("GET", ""));
+
+        Assert.Equal(HttpStatusCode.OK, result.StatusCode);
+
+        using var zipArchive = new ZipArchive(result.Body);
+        foreach (var entry in zipArchive.Entries) {
+            Assert.True(int.TryParse(entry.Name, out var index));
+            Assert.True(index >= 0 && index < files.Count);
+            using var entryStream = entry.Open();
+            using var sr = new StreamReader(entryStream);
+            var actualContent = sr.ReadToEnd();
+            Assert.Equal(files[index].content.ToString(), actualContent);
+        }
+    }
+}

src/cli/onefuzz/api.py

@@ -29,6 +29,7 @@ from onefuzztypes import (
 )
 from onefuzztypes.enums import TaskType
 from pydantic import BaseModel
+from requests import Response
 from six.moves import input  # workaround for static analysis
 
 from .__version__ import __version__
@@ -94,6 +95,23 @@ class Endpoint:
         self.onefuzz = onefuzz
         self.logger = onefuzz.logger
 
+    def _req_base(
+        self,
+        method: str,
+        *,
+        data: Optional[BaseModel] = None,
+        as_params: bool = False,
+        alternate_endpoint: Optional[str] = None,
+    ) -> Response:
+        endpoint = self.endpoint if alternate_endpoint is None else alternate_endpoint
+
+        if as_params:
+            response = self.onefuzz._backend.request(method, endpoint, params=data)
+        else:
+            response = self.onefuzz._backend.request(method, endpoint, json_data=data)
+
+        return response
+
     def _req_model(
         self,
         method: str,
@@ -103,12 +121,12 @@ class Endpoint:
         as_params: bool = False,
         alternate_endpoint: Optional[str] = None,
     ) -> A:
-        endpoint = self.endpoint if alternate_endpoint is None else alternate_endpoint
-
-        if as_params:
-            response = self.onefuzz._backend.request(method, endpoint, params=data)
-        else:
-            response = self.onefuzz._backend.request(method, endpoint, json_data=data)
+        response = self._req_base(
+            method,
+            data=data,
+            as_params=as_params,
+            alternate_endpoint=alternate_endpoint,
+        ).json()
 
         return model.parse_obj(response)
 
@@ -124,9 +142,13 @@ class Endpoint:
         endpoint = self.endpoint if alternate_endpoint is None else alternate_endpoint
 
         if as_params:
-            response = self.onefuzz._backend.request(method, endpoint, params=data)
+            response = self.onefuzz._backend.request(
+                method, endpoint, params=data
+            ).json()
         else:
-            response = self.onefuzz._backend.request(method, endpoint, json_data=data)
+            response = self.onefuzz._backend.request(
+                method, endpoint, json_data=data
+            ).json()
 
         return [model.parse_obj(x) for x in response]
 
@@ -1115,6 +1137,22 @@ class JobTasks(Endpoint):
         return self.onefuzz.tasks.list(job_id=job_id, state=[])
 
 
+class Tools(Endpoint):
+    """Interact with tasks within a job"""
+
+    endpoint = "tools"
+
+    def get(self, destination: str) -> str:
+        """Download a zip file containing the agent binaries"""
+        self.logger.debug("get tools")
+
+        response = self._req_base("GET")
+        path = os.path.join(destination, "tools.zip")
+        open(path, "wb").write(response.content)
+
+        return path
+
+
 class Jobs(Endpoint):
     """Interact with Jobs"""
 
@@ -1720,6 +1758,7 @@ class Onefuzz:
         self.scalesets = Scaleset(self)
         self.nodes = Node(self)
         self.webhooks = Webhooks(self)
+        self.tools = Tools(self)
         self.instance_config = InstanceConfigCmd(self)
 
         if self._backend.is_feature_enabled(PreviewFeature.job_templates.name):

src/cli/onefuzz/backend.py

@@ -32,6 +32,7 @@ import msal
 import requests
 from azure.storage.blob import ContainerClient
 from pydantic import BaseModel, Field
+from requests import Response
 from tenacity import RetryCallState, retry
 from tenacity.retry import retry_if_exception_type
 from tenacity.stop import stop_after_attempt
@@ -285,7 +286,7 @@ class Backend:
         json_data: Optional[Any] = None,
         params: Optional[Any] = None,
         _retry_on_auth_failure: bool = True,
-    ) -> Any:
+    ) -> Response:
         if self.config.dotnet_functions and path in self.config.dotnet_functions:
             endpoint = self.config.dotnet_endpoint
         else:
@@ -349,7 +350,7 @@ class Backend:
                 "request did not succeed: HTTP %s - %s"
                 % (response.status_code, error_text)
             )
-        return response.json()
+        return response
 
 
 def before_sleep(retry_state: RetryCallState) -> None: