Add extensions functionality for C# refactor (#2218)

* Add extensions functionality

* Small cleanup

* Use context, fix typo

src/ApiService/ApiService/Program.cs

@@ -109,6 +109,7 @@ public class Program {
             .AddScoped<INodeTasksOperations, NodeTasksOperations>()
             .AddScoped<INodeMessageOperations, NodeMessageOperations>()
             .AddScoped<IRequestHandling, RequestHandling>()
+            .AddScoped<IImageOperations, ImageOperations>()
             .AddScoped<IOnefuzzContext, OnefuzzContext>()
             .AddScoped<IEndpointAuthorization, EndpointAuthorization>()
             .AddScoped<INodeMessageOperations, NodeMessageOperations>()

src/ApiService/ApiService/onefuzzlib/Extension.cs

@@ -1,4 +1,6 @@
 using System.Text.Json;
+using System.Threading.Tasks;
+using Azure.Core;
 using Azure.ResourceManager.Compute;
 using Azure.Storage.Sas;
 using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
@@ -7,37 +9,32 @@ namespace Microsoft.OneFuzz.Service;
 
 public interface IExtensions {
     public Async.Task<IList<VirtualMachineScaleSetExtensionData>> FuzzExtensions(Pool pool, Scaleset scaleset);
+
+    public Async.Task<Dictionary<string, VirtualMachineExtensionData>> ReproExtensions(AzureLocation region, Os reproOs, Guid reproId, ReproConfig reproConfig, Container? setupContainer);
 }
 
-
 public class Extensions : IExtensions {
-    IServiceConfig _serviceConfig;
+    IOnefuzzContext _context;
-    ICreds _creds;
-    IQueue _queue;
-    IContainers _containers;
-    IConfigOperations _instanceConfigOps;
-    ILogAnalytics _logAnalytics;
 
-    public Extensions(IServiceConfig config, ICreds creds, IQueue queue, IContainers containers, IConfigOperations instanceConfigOps, ILogAnalytics logAnalytics) {
+    private static readonly JsonSerializerOptions _extensionSerializerOptions = new JsonSerializerOptions {
-        _serviceConfig = config;
+        PropertyNamingPolicy = JsonNamingPolicy.CamelCase
-        _creds = creds;
+    };
-        _queue = queue;
+
-        _containers = containers;
+    public Extensions(IOnefuzzContext context) {
-        _instanceConfigOps = instanceConfigOps;
+        _context = context;
-        _logAnalytics = logAnalytics;
     }
 
     public async Async.Task<Uri?> ConfigUrl(Container container, string fileName, bool withSas) {
         if (withSas)
-            return await _containers.GetFileSasUrl(container, fileName, StorageType.Config, BlobSasPermissions.Read);
+            return await _context.Containers.GetFileSasUrl(container, fileName, StorageType.Config, BlobSasPermissions.Read);
         else
-            return await _containers.GetFileUrl(container, fileName, StorageType.Config);
+            return await _context.Containers.GetFileUrl(container, fileName, StorageType.Config);
     }
 
-    public async Async.Task<IList<VirtualMachineScaleSetExtensionData>> GenericExtensions(string region, Os vmOs) {
+    public async Async.Task<IList<VMExtensionWrapper>> GenericExtensions(AzureLocation region, Os vmOs) {
-        var extensions = new List<VirtualMachineScaleSetExtensionData>();
+        var extensions = new List<VMExtensionWrapper>();
 
-        var instanceConfig = await _instanceConfigOps.Fetch();
+        var instanceConfig = await _context.ConfigOperations.Fetch();
         extensions.Add(await MonitorExtension(region, vmOs));
 
         var depenency = DependencyExtension(region, vmOs);
@@ -71,19 +68,20 @@ public class Extensions : IExtensions {
         return extensions;
     }
 
-    public static VirtualMachineScaleSetExtensionData KeyVaultExtension(string region, KeyvaultExtensionConfig keyVault, Os vmOs) {
+    public static VMExtensionWrapper KeyVaultExtension(AzureLocation region, KeyvaultExtensionConfig keyVault, Os vmOs) {
         var keyVaultName = keyVault.KeyVaultName;
         var certName = keyVault.CertName;
         var uri = keyVaultName + certName;
 
         if (vmOs == Os.Windows) {
-            return new VirtualMachineScaleSetExtensionData {
+            return new VMExtensionWrapper {
+                Location = region,
                 Name = "KVVMExtensionForWindows",
                 Publisher = "Microsoft.Azure.KeyVault",
                 TypePropertiesType = "KeyVaultForWindows",
                 TypeHandlerVersion = "1.0",
                 AutoUpgradeMinorVersion = true,
-                Settings = new BinaryData(new {
+                Settings = new BinaryData(JsonSerializer.Serialize(new {
                     SecretsManagementSettings = new {
                         PollingIntervalInS = "3600",
                         CertificateStoreName = "MY",
@@ -92,46 +90,48 @@ public class Extensions : IExtensions {
                         RequireInitialSync = true,
                         ObservedCertificates = new string[] { uri },
                     }
-                })
+                }, _extensionSerializerOptions))
             };
         } else if (vmOs == Os.Linux) {
             var certPath = keyVault.CertPath;
             var extensionStore = keyVault.ExtensionStore;
             var certLocation = certPath + extensionStore;
 
-            return new VirtualMachineScaleSetExtensionData {
+            return new VMExtensionWrapper {
+                Location = region,
                 Name = "KVVMExtensionForLinux",
                 Publisher = "Microsoft.Azure.KeyVault",
                 TypePropertiesType = "KeyVaultForLinux",
                 TypeHandlerVersion = "2.0",
                 AutoUpgradeMinorVersion = true,
-                Settings = new BinaryData(new {
+                Settings = new BinaryData(JsonSerializer.Serialize(new {
                     SecretsManagementSettings = new {
                         PollingIntervalInS = "3600",
                         CertificateStoreLocation = certLocation,
                         RequireInitialSync = true,
                         ObservedCertificates = new string[] { uri },
                     }
-                })
+                }, _extensionSerializerOptions))
             };
         } else {
             throw new NotImplementedException($"unsupported os {vmOs}");
         }
     }
 
-    public static VirtualMachineScaleSetExtensionData AzSecExtension(string region) {
+    public static VMExtensionWrapper AzSecExtension(AzureLocation region) {
-        return new VirtualMachineScaleSetExtensionData {
+        return new VMExtensionWrapper {
+            Location = region,
             Name = "AzureSecurityLinuxAgent",
             Publisher = "Microsoft.Azure.Security.Monitoring",
             TypePropertiesType = "AzureSecurityLinuxAgent",
             TypeHandlerVersion = "2.0",
             AutoUpgradeMinorVersion = true,
-            Settings = new BinaryData(new { EnableGenevaUpload = true, EnableAutoConfig = true })
+            Settings = new BinaryData(JsonSerializer.Serialize(new { EnableGenevaUpload = true, EnableAutoConfig = true }, _extensionSerializerOptions))
         };
 
     }
 
-    public static VirtualMachineScaleSetExtensionData AzMonExtension(string region, AzureMonitorExtensionConfig azureMonitor) {
+    public static VMExtensionWrapper AzMonExtension(AzureLocation region, AzureMonitorExtensionConfig azureMonitor) {
         var authId = azureMonitor.MonitoringGCSAuthId;
         var configVersion = azureMonitor.ConfigVersion;
         var moniker = azureMonitor.Moniker;
@@ -140,15 +140,16 @@ public class Extensions : IExtensions {
         var account = azureMonitor.MonitoringGCSAccount;
         var authIdType = azureMonitor.MonitoringGCSAuthIdType;
 
-        return new VirtualMachineScaleSetExtensionData {
+        return new VMExtensionWrapper {
+            Location = region,
             Name = "AzureMonitorLinuxAgent",
             Publisher = "Microsoft.Azure.Monitor",
             TypePropertiesType = "AzureMonitorLinuxAgent",
             AutoUpgradeMinorVersion = true,
             TypeHandlerVersion = "1.0",
-            Settings = new BinaryData(new { GCS_AUTO_CONFIG = true }),
+            Settings = new BinaryData(JsonSerializer.Serialize(new { GCS_AUTO_CONFIG = true }, _extensionSerializerOptions)),
             ProtectedSettings =
-                new BinaryData(
+                new BinaryData(JsonSerializer.Serialize(
                     new {
                         ConfigVersion = configVersion,
                         Moniker = moniker,
@@ -158,14 +159,13 @@ public class Extensions : IExtensions {
                         MonitoringGCSRegion = region,
                         MonitoringGCSAuthId = authId,
                         MonitoringGCSAuthIdType = authIdType,
-                    })
+                    }, _extensionSerializerOptions))
         };
     }
 
-
+    public static VMExtensionWrapper GenevaExtension(AzureLocation region) {
-
+        return new VMExtensionWrapper {
-    public static VirtualMachineScaleSetExtensionData GenevaExtension(string region) {
+            Location = region,
-        return new VirtualMachineScaleSetExtensionData {
             Name = "Microsoft.Azure.Geneva.GenevaMonitoring",
             Publisher = "Microsoft.Azure.Geneva",
             TypePropertiesType = "GenevaMonitoring",
@@ -175,12 +175,13 @@ public class Extensions : IExtensions {
         };
     }
 
-    public static VirtualMachineScaleSetExtensionData? DependencyExtension(string region, Os vmOs) {
+    public static VMExtensionWrapper? DependencyExtension(AzureLocation region, Os vmOs) {
 
         if (vmOs == Os.Windows) {
-            return new VirtualMachineScaleSetExtensionData {
+            return new VMExtensionWrapper {
-                AutoUpgradeMinorVersion = true,
+                Location = region,
                 Name = "DependencyAgentWindows",
+                AutoUpgradeMinorVersion = true,
                 Publisher = "Microsoft.Azure.Monitoring.DependencyAgent",
                 TypePropertiesType = "DependencyAgentWindows",
                 TypeHandlerVersion = "9.5"
@@ -202,22 +203,22 @@ public class Extensions : IExtensions {
 
 
     public async Async.Task<Uri?> BuildPoolConfig(Pool pool) {
-        var instanceId = await _containers.GetInstanceId();
+        var instanceId = await _context.Containers.GetInstanceId();
 
-        var queueSas = await _queue.GetQueueSas("node-heartbeat", StorageType.Config, QueueSasPermissions.Add);
+        var queueSas = await _context.Queue.GetQueueSas("node-heartbeat", StorageType.Config, QueueSasPermissions.Add);
         var config = new AgentConfig(
             ClientCredentials: null,
-            OneFuzzUrl: _creds.GetInstanceUrl(),
+            OneFuzzUrl: _context.Creds.GetInstanceUrl(),
             PoolName: pool.Name,
             HeartbeatQueue: queueSas,
-            InstanceTelemetryKey: _serviceConfig.ApplicationInsightsInstrumentationKey,
+            InstanceTelemetryKey: _context.ServiceConfiguration.ApplicationInsightsInstrumentationKey,
-            MicrosoftTelemetryKey: _serviceConfig.OneFuzzTelemetry,
+            MicrosoftTelemetryKey: _context.ServiceConfiguration.OneFuzzTelemetry,
-            MultiTenantDomain: _serviceConfig.MultiTenantDomain,
+            MultiTenantDomain: _context.ServiceConfiguration.MultiTenantDomain,
             InstanceId: instanceId
             );
 
         var fileName = $"{pool.Name}/config.json";
-        await _containers.SaveBlob(new Container("vm-scripts"), fileName, (JsonSerializer.Serialize(config, EntityConverter.GetJsonSerializerOptions())), StorageType.Config);
+        await _context.Containers.SaveBlob(new Container("vm-scripts"), fileName, (JsonSerializer.Serialize(config, EntityConverter.GetJsonSerializerOptions())), StorageType.Config);
         return await ConfigUrl(new Container("vm-scripts"), fileName, false);
     }
 
@@ -234,25 +235,24 @@ public class Extensions : IExtensions {
             commands.Add($"Set-Content -Path {sshPath} -Value \"{sshKey}\"");
         }
 
-        await _containers.SaveBlob(new Container("vm-scripts"), fileName, string.Join(sep, commands) + sep, StorageType.Config);
+        await _context.Containers.SaveBlob(new Container("vm-scripts"), fileName, string.Join(sep, commands) + sep, StorageType.Config);
-        return await _containers.GetFileUrl(new Container("vm-scripts"), fileName, StorageType.Config);
+        return await _context.Containers.GetFileUrl(new Container("vm-scripts"), fileName, StorageType.Config);
     }
 
     public async Async.Task UpdateManagedScripts() {
-        var instanceSpecificSetupSas = _containers.GetContainerSasUrl(new Container("instance-specific-setup"), StorageType.Config, BlobContainerSasPermissions.List | BlobContainerSasPermissions.Read);
+        var instanceSpecificSetupSas = await _context.Containers.GetContainerSasUrl(new Container("instance-specific-setup"), StorageType.Config, BlobContainerSasPermissions.List | BlobContainerSasPermissions.Read);
-        var toolsSas = _containers.GetContainerSasUrl(new Container("tools"), StorageType.Config, BlobContainerSasPermissions.List | BlobContainerSasPermissions.Read);
+        var toolsSas = await _context.Containers.GetContainerSasUrl(new Container("tools"), StorageType.Config, BlobContainerSasPermissions.List | BlobContainerSasPermissions.Read);
 
         string[] commands = {
             $"azcopy sync '{instanceSpecificSetupSas}' instance-specific-setup",
             $"azcopy sync '{toolsSas}' tools"
         };
 
-        await _containers.SaveBlob(new Container("vm-scripts"), "managed.ps1", string.Join("\r\n", commands) + "\r\n", StorageType.Config);
+        await _context.Containers.SaveBlob(new Container("vm-scripts"), "managed.ps1", string.Join("\r\n", commands) + "\r\n", StorageType.Config);
-        await _containers.SaveBlob(new Container("vm-scripts"), "managed.sh", string.Join("\n", commands) + "\n", StorageType.Config);
+        await _context.Containers.SaveBlob(new Container("vm-scripts"), "managed.sh", string.Join("\n", commands) + "\n", StorageType.Config);
     }
 
-
+    public async Async.Task<VMExtensionWrapper> AgentConfig(AzureLocation region, Os vmOs, AgentMode mode, List<Uri>? urls = null, bool withSas = false) {
-    public async Async.Task<VirtualMachineScaleSetExtensionData> AgentConfig(string region, Os vmOs, AgentMode mode, List<Uri>? urls = null, bool withSas = false) {
         await UpdateManagedScripts();
         var urlsUpdated = urls ?? new();
 
@@ -267,17 +267,18 @@ public class Extensions : IExtensions {
             urlsUpdated.Add(toolsSetup);
             urlsUpdated.Add(toolsOneFuzz);
 
-            var toExecuteCmd = $"powershell -ExecutionPolicy Unrestricted -File win64/setup.ps1 -mode {mode}";
+            var toExecuteCmd = $"powershell -ExecutionPolicy Unrestricted -File win64/setup.ps1 -mode {mode.ToString().ToLowerInvariant()}";
 
-            var extension = new VirtualMachineScaleSetExtensionData {
+            var extension = new VMExtensionWrapper {
                 Name = "CustomScriptExtension",
                 TypePropertiesType = "CustomScriptExtension",
                 Publisher = "Microsoft.Compute",
+                Location = region,
                 ForceUpdateTag = Guid.NewGuid().ToString(),
                 TypeHandlerVersion = "1.9",
                 AutoUpgradeMinorVersion = true,
-                Settings = new BinaryData(new { commandToExecute = toExecuteCmd, fileUrls = urlsUpdated }),
+                Settings = new BinaryData(JsonSerializer.Serialize(new { commandToExecute = toExecuteCmd, fileUris = urlsUpdated }, _extensionSerializerOptions)),
-                ProtectedSettings = new BinaryData(new { managedIdentity = new Dictionary<string, string>() })
+                ProtectedSettings = new BinaryData(JsonSerializer.Serialize(new { managedIdentity = new Dictionary<string, string>() }, _extensionSerializerOptions))
             };
             return extension;
         } else if (vmOs == Os.Linux) {
@@ -290,17 +291,20 @@ public class Extensions : IExtensions {
             urlsUpdated.Add(toolsAzCopy);
             urlsUpdated.Add(toolsSetup);
 
-            var toExecuteCmd = $"sh setup.sh {mode}";
+            var toExecuteCmd = $"sh setup.sh {mode.ToString().ToLowerInvariant()}";
+            var extensionSettings = JsonSerializer.Serialize(new { CommandToExecute = toExecuteCmd, FileUris = urlsUpdated }, _extensionSerializerOptions);
+            var protectedExtensionSettings = JsonSerializer.Serialize(new { ManagedIdentity = new Dictionary<string, string>() }, _extensionSerializerOptions);
 
-            var extension = new VirtualMachineScaleSetExtensionData {
+            var extension = new VMExtensionWrapper {
                 Name = "CustomScript",
+                Publisher = "Microsoft.Azure.Extensions",
                 TypePropertiesType = "CustomScript",
-                Publisher = "Microsoft.Azure.Extension",
-                ForceUpdateTag = Guid.NewGuid().ToString(),
                 TypeHandlerVersion = "2.1",
+                Location = region,
+                ForceUpdateTag = Guid.NewGuid().ToString(),
                 AutoUpgradeMinorVersion = true,
-                Settings = new BinaryData(new { CommandToExecute = toExecuteCmd, FileUrls = urlsUpdated }),
+                Settings = new BinaryData(extensionSettings),
-                ProtectedSettings = new BinaryData(new { ManagedIdentity = new Dictionary<string, string>() })
+                ProtectedSettings = new BinaryData(protectedExtensionSettings)
             };
             return extension;
         }
@@ -308,28 +312,31 @@ public class Extensions : IExtensions {
         throw new NotImplementedException($"unsupported OS: {vmOs}");
     }
 
-    public async Async.Task<VirtualMachineScaleSetExtensionData> MonitorExtension(string region, Os vmOs) {
+    public async Async.Task<VMExtensionWrapper> MonitorExtension(AzureLocation region, Os vmOs) {
-        var settings = await _logAnalytics.GetMonitorSettings();
+        var settings = await _context.LogAnalytics.GetMonitorSettings();
-
+        var extensionSettings = JsonSerializer.Serialize(new { WorkspaceId = settings.Id }, _extensionSerializerOptions);
+        var protectedExtensionSettings = JsonSerializer.Serialize(new { WorkspaceKey = settings.Key }, _extensionSerializerOptions);
         if (vmOs == Os.Windows) {
-            return new VirtualMachineScaleSetExtensionData {
+            return new VMExtensionWrapper {
+                Location = region,
                 Name = "OMSExtension",
                 TypePropertiesType = "MicrosoftMonitoringAgent",
                 Publisher = "Microsoft.EnterpriseCloud.Monitoring",
                 TypeHandlerVersion = "1.0",
                 AutoUpgradeMinorVersion = true,
-                Settings = new BinaryData(new { WorkSpaceId = settings.Id }),
+                Settings = new BinaryData(extensionSettings),
-                ProtectedSettings = new BinaryData(new { WorkspaceKey = settings.Key })
+                ProtectedSettings = new BinaryData(protectedExtensionSettings)
             };
         } else if (vmOs == Os.Linux) {
-            return new VirtualMachineScaleSetExtensionData {
+            return new VMExtensionWrapper {
+                Location = region,
                 Name = "OMSExtension",
                 TypePropertiesType = "OmsAgentForLinux",
                 Publisher = "Microsoft.EnterpriseCloud.Monitoring",
                 TypeHandlerVersion = "1.12",
                 AutoUpgradeMinorVersion = true,
-                Settings = new BinaryData(new { WorkSpaceId = settings.Id }),
+                Settings = new BinaryData(extensionSettings),
-                ProtectedSettings = new BinaryData(new { WorkspaceKey = settings.Key })
+                ProtectedSettings = new BinaryData(protectedExtensionSettings)
             };
         } else {
             throw new NotImplementedException($"unsupported os: {vmOs}");
@@ -346,6 +353,100 @@ public class Extensions : IExtensions {
         var extensions = await GenericExtensions(scaleset.Region, pool.Os);
 
         extensions.Add(fuzzExtension);
-        return extensions;
+        return extensions.Select(extension => extension.GetAsVirtualMachineScaleSetExtension()).ToList();
     }
+
+    public async Task<Dictionary<string, VirtualMachineExtensionData>> ReproExtensions(AzureLocation region, Os reproOs, Guid reproId, ReproConfig reproConfig, Container? setupContainer) {
+        // TODO: what about contents of repro.ps1 / repro.sh?
+        var report = await _context.Reports.GetReport(reproConfig.Container, reproConfig.Path);
+        report.EnsureNotNull($"invalid report: {reproConfig}");
+        report?.InputBlob.EnsureNotNull("unable to perform reproduction without an input blob");
+
+        var commands = new List<string>();
+        if (setupContainer != null) {
+            var containerSasUrl = await _context.Containers.GetContainerSasUrl(
+                setupContainer,
+                StorageType.Corpus,
+                BlobContainerSasPermissions.Read | BlobContainerSasPermissions.List
+            );
+            commands.Add(
+                $"azcopy sync '{containerSasUrl}' ./setup"
+            );
+        }
+
+        var urls = new List<Uri>()
+        {
+            await _context.Containers.GetFileSasUrl(
+                reproConfig.Container,
+                reproConfig.Path,
+                StorageType.Corpus,
+                BlobSasPermissions.Read
+            ),
+            await _context.Containers.GetFileSasUrl(
+                report?.InputBlob?.container!,
+                report?.InputBlob?.Name!,
+                StorageType.Corpus,
+                BlobSasPermissions.Read
+            )
+        };
+
+        List<string> reproFiles;
+        string taskScript;
+        string scriptName;
+        if (reproOs == Os.Windows) {
+            reproFiles = new List<string>()
+            {
+                $"{reproId}/repro.ps1"
+            };
+            taskScript = string.Join("\r\n", commands);
+            scriptName = "task-setup.ps1";
+        } else {
+            reproFiles = new List<string>()
+            {
+                $"{reproId}/repro.sh",
+                $"{reproId}/repro-stdout.sh"
+            };
+            commands.Add("chmod -R +x setup");
+            taskScript = string.Join("\n", commands);
+            scriptName = "task-setup.sh";
+        }
+
+        await _context.Containers.SaveBlob(
+            new Container("task-configs"),
+            $"{reproId}/{scriptName}",
+            taskScript,
+            StorageType.Config
+        );
+
+        foreach (var reproFile in reproFiles) {
+            urls.AddRange(new List<Uri>()
+            {
+                await _context.Containers.GetFileSasUrl(
+                    new Container("repro-scripts"),
+                    reproFile,
+                    StorageType.Config,
+                    BlobSasPermissions.Read
+                ),
+                await _context.Containers.GetFileSasUrl(
+                    new Container("task-configs"),
+                    $"{reproId}/{scriptName}",
+                    StorageType.Config,
+                    BlobSasPermissions.Read
+                )
+            });
+        }
+
+        var baseExtension = await AgentConfig(region, reproOs, AgentMode.Repro, urls: urls, withSas: true);
+        var extensions = await GenericExtensions(region, reproOs);
+        extensions.Add(baseExtension);
+
+        var extensionsDict = new Dictionary<string, VirtualMachineExtensionData>();
+        foreach (var extension in extensions) {
+            var (name, data) = extension.GetAsVirtualMachineExtension();
+            extensionsDict.Add(name, data);
+        }
+
+        return extensionsDict;
+    }
+
 }

src/ApiService/ApiService/onefuzzlib/ImageOperations.cs

@@ -0,0 +1,115 @@
+using System.Threading.Tasks;
+using Azure;
+using Azure.ResourceManager.Compute;
+
+namespace Microsoft.OneFuzz.Service;
+
+public interface IImageOperations {
+    public Async.Task<OneFuzzResult<Os>> GetOs(string region, string image);
+}
+
+public class ImageOperations : IImageOperations {
+    private IOnefuzzContext _context;
+    private ILogTracer _logTracer;
+
+    public ImageOperations(ILogTracer logTracer, IOnefuzzContext context) {
+        _logTracer = logTracer;
+        _context = context;
+    }
+    public async Task<OneFuzzResult<Os>> GetOs(string region, string image) {
+        string? name = null;
+        try {
+            var parsed = _context.Creds.ParseResourceId(image);
+            parsed = await _context.Creds.GetData(parsed);
+            if (string.Equals(parsed.Id.ResourceType, "galleries", StringComparison.OrdinalIgnoreCase)) {
+                try {
+                    // This is not _exactly_ the same as the python code  
+                    // because in C# we don't have access to child_name_1
+                    var gallery = await _context.Creds.GetResourceGroupResource().GetGalleries().GetAsync(
+                        parsed.Data.Name
+                    );
+
+                    var galleryImage = gallery.Value.GetGalleryImages()
+                        .ToEnumerable()
+                        .Where(galleryImage => string.Equals(galleryImage.Id, parsed.Id, StringComparison.OrdinalIgnoreCase))
+                        .First();
+
+                    galleryImage = await galleryImage.GetAsync();
+
+                    name = galleryImage.Data?.OSType?.ToString().ToLowerInvariant()!;
+
+                } catch (Exception ex) when (
+                      ex is RequestFailedException ||
+                      ex is NullReferenceException
+                  ) {
+                    return OneFuzzResult<Os>.Error(
+                        ErrorCode.INVALID_IMAGE,
+                        ex.ToString()
+                    );
+                }
+            } else {
+                try {
+                    name = (await _context.Creds.GetResourceGroupResource().GetImages().GetAsync(
+                        parsed.Data.Name
+                    )).Value.Data.StorageProfile.OSDisk.OSType.ToString().ToLowerInvariant();
+                } catch (Exception ex) when (
+                    ex is RequestFailedException ||
+                    ex is NullReferenceException
+                ) {
+                    return OneFuzzResult<Os>.Error(
+                        ErrorCode.INVALID_IMAGE,
+                        ex.ToString()
+                    );
+                }
+            }
+        } catch (FormatException) {
+            var imageParts = image.Split(":");
+
+            // The python code would throw if more than 4 parts are found in the split
+            System.Diagnostics.Trace.Assert(imageParts.Length == 4, $"Expected 4 ':' separated parts in {image}");
+
+            var publisher = imageParts[0];
+            var offer = imageParts[1];
+            var sku = imageParts[2];
+            var version = imageParts[3];
+
+            try {
+                var subscription = await _context.Creds.ArmClient.GetDefaultSubscriptionAsync();
+                if (string.Equals(version, "latest", StringComparison.Ordinal)) {
+                    version = (await subscription.GetVirtualMachineImagesAsync(
+                        region,
+                        publisher,
+                        offer,
+                        sku,
+                        top: 1
+                    ).FirstAsync()).Name;
+                }
+
+                name = (await subscription.GetVirtualMachineImageAsync(
+                    region,
+                    publisher,
+                    offer,
+                    sku
+                    , version
+                )).Value.OSDiskImageOperatingSystem.ToString().ToLower();
+            } catch (RequestFailedException ex) {
+                return OneFuzzResult<Os>.Error(
+                    ErrorCode.INVALID_IMAGE,
+                    ex.ToString()
+                );
+            }
+        }
+
+        if (name != null) {
+            name = string.Concat(name[0].ToString().ToUpper(), name.AsSpan(1));
+            if (Enum.TryParse(name, out Os os)) {
+                return OneFuzzResult<Os>.Ok(os);
+            }
+        }
+
+        return OneFuzzResult<Os>.Error(
+            ErrorCode.INVALID_IMAGE,
+            $"Unexpected image os type: {name}"
+        );
+    }
+}

src/ApiService/ApiService/onefuzzlib/OnefuzzContext.cs

@@ -38,6 +38,7 @@ public interface IOnefuzzContext {
     IRequestHandling RequestHandling { get; }
     INsgOperations NsgOperations { get; }
     ISubnet Subnet { get; }
+    IImageOperations ImageOperations { get; }
 }
 
 public class OnefuzzContext : IOnefuzzContext {
@@ -81,4 +82,5 @@ public class OnefuzzContext : IOnefuzzContext {
     public IRequestHandling RequestHandling => _serviceProvider.GetRequiredService<IRequestHandling>();
     public INsgOperations NsgOperations => _serviceProvider.GetRequiredService<INsgOperations>();
     public ISubnet Subnet => _serviceProvider.GetRequiredService<ISubnet>();
+    public IImageOperations ImageOperations => _serviceProvider.GetRequiredService<IImageOperations>();
 }

src/ApiService/ApiService/onefuzzlib/Reports.cs

@@ -1,10 +1,12 @@
 using System.Text.Json;
+using System.Threading.Tasks;
 using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
 
 namespace Microsoft.OneFuzz.Service;
 
 public interface IReports {
     public Async.Task<IReport?> GetReportOrRegression(Container container, string fileName, bool expectReports = false, params string[] args);
+    public Async.Task<Report?> GetReport(Container container, string fileName);
 }
 
 public class Reports : IReports {
@@ -15,6 +17,15 @@ public class Reports : IReports {
         _containers = containers;
     }
 
+    public async Task<Report?> GetReport(Container container, string fileName) {
+        var result = await GetReportOrRegression(container, fileName);
+        if (result != null && result is Report) {
+            return result as Report;
+        }
+
+        return null;
+    }
+
     public async Async.Task<IReport?> GetReportOrRegression(Container container, string fileName, bool expectReports = false, params string[] args) {
         var filePath = String.Join("/", new[] { container.ContainerName, fileName });
         if (!fileName.EndsWith(".json", StringComparison.Ordinal)) {

src/ApiService/ApiService/onefuzzlib/VmExtensionWrapper.cs

@@ -0,0 +1,62 @@
+using Azure.Core;
+using Azure.ResourceManager.Compute;
+
+namespace Microsoft.OneFuzz.Service {
+    public class VMExtensionWrapper {
+        public AzureLocation? Location { get; init; }
+        public string? Name { get; init; }
+        public string? TypePropertiesType { get; init; }
+        public string? Publisher { get; init; }
+        public string? TypeHandlerVersion { get; init; }
+        public string? ForceUpdateTag { get; init; }
+        public bool? AutoUpgradeMinorVersion { get; init; }
+        public bool? EnableAutomaticUpgrade { get; init; }
+        public BinaryData? Settings { get; init; }
+        public BinaryData? ProtectedSettings { get; init; }
+
+        public (string, VirtualMachineExtensionData) GetAsVirtualMachineExtension() {
+            if (Location == null) { // EnsureNotNull does not satisfy the nullability checker
+                throw new ArgumentNullException("Location required for VirtualMachineExtension");
+            }
+            TypePropertiesType.EnsureNotNull("TypePropertiesType required for VirtualMachineExtension");
+            Publisher.EnsureNotNull("Publisher required for VirtualMachineExtension");
+            TypeHandlerVersion.EnsureNotNull("TypeHandlerVersion required for VirtualMachineExtension");
+            AutoUpgradeMinorVersion.EnsureNotNull("AutoUpgradeMinorVersion required for VirtualMachineExtension");
+            Settings.EnsureNotNull("Settings required for VirtualMachineExtension");
+            ProtectedSettings.EnsureNotNull("ProtectedSettings required for VirtualMachineExtension");
+
+            return (Name!, new VirtualMachineExtensionData(Location.Value) {
+                TypePropertiesType = TypePropertiesType,
+                Publisher = Publisher,
+                TypeHandlerVersion = TypeHandlerVersion,
+                AutoUpgradeMinorVersion = AutoUpgradeMinorVersion,
+                EnableAutomaticUpgrade = EnableAutomaticUpgrade,
+                ForceUpdateTag = ForceUpdateTag,
+                Settings = Settings,
+                ProtectedSettings = ProtectedSettings
+            });
+        }
+
+        public VirtualMachineScaleSetExtensionData GetAsVirtualMachineScaleSetExtension() {
+            Name.EnsureNotNull("Name required for VirtualMachineScaleSetExtension");
+            TypePropertiesType.EnsureNotNull("TypePropertiesType required for VirtualMachineScaleSetExtension");
+            Publisher.EnsureNotNull("Publisher required for VirtualMachineScaleSetExtension");
+            TypeHandlerVersion.EnsureNotNull("TypeHandlerVersion required for VirtualMachineScaleSetExtension");
+            AutoUpgradeMinorVersion.EnsureNotNull("AutoUpgradeMinorVersion required for VirtualMachineScaleSetExtension");
+            Settings.EnsureNotNull("Settings required for VirtualMachineScaleSetExtension");
+            ProtectedSettings.EnsureNotNull("ProtectedSettings required for VirtualMachineScaleSetExtension");
+            return new VirtualMachineScaleSetExtensionData() {
+                Name = Name,
+                TypePropertiesType = TypePropertiesType,
+                Publisher = Publisher,
+                TypeHandlerVersion = TypeHandlerVersion,
+                AutoUpgradeMinorVersion = AutoUpgradeMinorVersion,
+                EnableAutomaticUpgrade = EnableAutomaticUpgrade,
+                ForceUpdateTag = ForceUpdateTag,
+                Settings = Settings,
+                ProtectedSettings = ProtectedSettings
+            };
+        }
+    }
+
+}

src/ApiService/ApiService/onefuzzlib/orm/Orm.cs

@@ -210,6 +210,8 @@ namespace ApiService.OneFuzzLib.Orm {
             if (func != null) {
                 _logTracer.Info($"processing state update: {typeof(T)} - PartitionKey {_partitionKeyGetter?.Value()} {_rowKeyGetter?.Value()} - %s");
                 return await func(entity);
+            } else {
+                _logTracer.Info($"State function for state: '{state}' not found on type {typeof(T)}");
             }
             return null;
         }

src/ApiService/IntegrationTests/Fakes/TestContext.cs

@@ -109,4 +109,6 @@ public sealed class TestContext : IOnefuzzContext {
     public INsgOperations NsgOperations => throw new NotImplementedException();
 
     public ISubnet Subnet => throw new NotImplementedException();
+
+    public IImageOperations ImageOperations => throw new NotImplementedException();
 }