From 1ade4bdf2d6ee70b0ed1c485265e93593d2dbd2e Mon Sep 17 00:00:00 2001
From: Kevin Anderson <keander@microsoft.com>
Date: Mon, 25 Nov 2024 15:02:40 -0800
Subject: [PATCH] Adding ability to define a custom AzureEnvironment.

---
 src/Promitor.Core/Extensions/AzureCloudExtensions.cs |  4 ++++
 .../Extensions/AzureEnvironmentExtensions.cs         | 12 ++++++++++++
 src/Promitor.Core/Serialization/Enum/AzureCloud.cs   |  3 ++-
 .../LogAnalyticsClient.cs                            |  2 ++
 4 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/src/Promitor.Core/Extensions/AzureCloudExtensions.cs b/src/Promitor.Core/Extensions/AzureCloudExtensions.cs
index 0980bd118..483590bb4 100644
--- a/src/Promitor.Core/Extensions/AzureCloudExtensions.cs
+++ b/src/Promitor.Core/Extensions/AzureCloudExtensions.cs
@@ -25,6 +25,8 @@ public static AzureEnvironment GetAzureEnvironment(this AzureCloud azureCloud)
                     return AzureEnvironment.AzureGermanCloud;
                 case AzureCloud.UsGov:
                     return AzureEnvironment.AzureUSGovernment;
+                case AzureCloud.Custom:
+                    return AzureEnvironmentExtensions.AzureCustomCloud;
                 default:
                     throw new ArgumentOutOfRangeException(nameof(azureCloud), "No Azure environment is known for in legacy SDK");
             }
@@ -81,6 +83,8 @@ public static Uri GetAzureAuthorityHost(this AzureCloud azureCloud)
                     return AzureAuthorityHosts.AzureGermany;
                 case AzureCloud.UsGov:
                     return AzureAuthorityHosts.AzureGovernment;
+                case AzureCloud.Custom:
+                    return new Uri(Environment.GetEnvironmentVariable("PROMITOR_AZURE_AUTH_ENDPOINT"));
                 default:
                     throw new ArgumentOutOfRangeException(nameof(azureCloud), "No Azure environment is known for");
             }
diff --git a/src/Promitor.Core/Extensions/AzureEnvironmentExtensions.cs b/src/Promitor.Core/Extensions/AzureEnvironmentExtensions.cs
index 5a0f2682c..8d5861c66 100644
--- a/src/Promitor.Core/Extensions/AzureEnvironmentExtensions.cs
+++ b/src/Promitor.Core/Extensions/AzureEnvironmentExtensions.cs
@@ -1,5 +1,6 @@
 using Humanizer;
 using Microsoft.Azure.Management.ResourceManager.Fluent;
+using System;
 
 namespace Promitor.Core.Extensions
 {
@@ -14,5 +15,16 @@ public static string GetDisplayName(this AzureEnvironment azureCloud)
         {
             return azureCloud.Name.Replace("Azure", "").Replace("Cloud", "").Humanize(LetterCasing.Title);
         }
+
+        public static AzureEnvironment AzureCustomCloud = new AzureEnvironment()
+        {
+            Name = nameof(AzureCustomCloud),
+            AuthenticationEndpoint = Environment.GetEnvironmentVariable("PROMITOR_AUTH_ENDPOINT"),
+            ResourceManagerEndpoint = Environment.GetEnvironmentVariable("PROMITOR_RESOURCE_MANAGER_ENDPOINT"),
+            ManagementEndpoint = Environment.GetEnvironmentVariable("PROMITOR_MANAGEMENT_ENDPOINT"),
+            GraphEndpoint = Environment.GetEnvironmentVariable("PROMITOR_GRAPH_ENDPOINT"),
+            StorageEndpointSuffix = Environment.GetEnvironmentVariable("PROMITOR_STORAGE_ENDPOINT_SUFFIX"),
+            KeyVaultSuffix = Environment.GetEnvironmentVariable("PROMITOR_KEY_VAULT_SUFFIX")
+        };
     }
 }
\ No newline at end of file
diff --git a/src/Promitor.Core/Serialization/Enum/AzureCloud.cs b/src/Promitor.Core/Serialization/Enum/AzureCloud.cs
index 32a663e2b..0be43fc6a 100644
--- a/src/Promitor.Core/Serialization/Enum/AzureCloud.cs
+++ b/src/Promitor.Core/Serialization/Enum/AzureCloud.cs
@@ -6,6 +6,7 @@ public enum AzureCloud
         Global,
         China,
         UsGov,
-        Germany
+        Germany,
+        Custom
     }
 }
\ No newline at end of file
diff --git a/src/Promitor.Integrations.LogAnalytics/LogAnalyticsClient.cs b/src/Promitor.Integrations.LogAnalytics/LogAnalyticsClient.cs
index c4f921a27..9e605f931 100644
--- a/src/Promitor.Integrations.LogAnalytics/LogAnalyticsClient.cs
+++ b/src/Promitor.Integrations.LogAnalytics/LogAnalyticsClient.cs
@@ -15,6 +15,7 @@ public class LogAnalyticsClient
         public readonly string ColumnNameResult = "result";
         private readonly Uri _defaultEndpoint = new("https://api.loganalytics.io");
         private readonly Uri _govEndpoint = new("https://api.loganalytics.us");
+        private readonly Uri _customEndpoint = new (Environment.GetEnvironmentVariable("PROMITOR_LOG_ANALYTICS_ENDPOINT"));
 
         private readonly LogsQueryClient _logsQueryClient;
 
@@ -31,6 +32,7 @@ public LogAnalyticsClient(ILoggerFactory loggerFactory, AzureEnvironment azureEn
             {
                 nameof(AzureEnvironment.AzureGlobalCloud) => _defaultEndpoint,
                 nameof(AzureEnvironment.AzureUSGovernment) => _govEndpoint,
+                "AzureCustomCloud" => _customEndpoint,
                 _ => throw new NotSupportedException($"Environment {azureEnvironment.Name} is not supported for scraping Azure Log Analytics resource(s)")
             };