Improve the interface

Azure · Aug 20, 2024 · b5b77e1 · b5b77e1
1 parent 8468613
commit b5b77e1
Show file tree

Hide file tree

Showing 18 changed files with 147 additions and 32 deletions.
diff --git a/src/Authentication.Abstractions/AuthenticationTelemetry.cs b/src/Authentication.Abstractions/AuthenticationTelemetry.cs
@@ -4,12 +4,6 @@ namespace Microsoft.Azure.Commands.Common.Authentication.Abstractions
 {
     public class AuthenticationTelemetry: IAzureTelemetry <AuthTelemetryRecord> 
     {
-        //public delegate string RequestIdAccquirer();
-
-        //public static RequestIdAccquirer GetRequestId;
-
-        //private static string requestId;
-
         public AuthenticationTelemetryData GetTelemetryRecord(ICmdletContext cmdletContext)
         {
             var records = PopTelemetryRecord(cmdletContext);

diff --git a/src/Authentication.Abstractions/Extensions/CmdletContextExtension.cs b/src/Authentication.Abstractions/Extensions/CmdletContextExtension.cs
@@ -0,0 +1,32 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using Microsoft.Azure.Commands.Common.Authentication.Abstractions.Interfaces;
+
+using System.Collections.Generic;
+
+namespace Microsoft.Azure.Commands.Common.Authentication.Abstractions.Extensions
+{
+    public static class CmdletContextExtension
+    {
+        public static IDictionary<string, object> ToExtensibleParameters(this ICmdletContext cmdletContext)
+        {
+            if (cmdletContext != null)
+            {
+                return new Dictionary<string, object> { { nameof(ICmdletContext), cmdletContext } };
+            }
+            return null;
+        }
+    }
+}
diff --git a/src/Authentication.Abstractions/Interfaces/IAuthenticationFactory.cs b/src/Authentication.Abstractions/Interfaces/IAuthenticationFactory.cs
@@ -16,6 +16,7 @@
 using System;
 using System.Security;
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions.Interfaces;
+using System.Collections.Generic;
 
 namespace Microsoft.Azure.Commands.Common.Authentication.Abstractions
 {
@@ -24,6 +25,27 @@ namespace Microsoft.Azure.Commands.Common.Authentication.Abstractions
     /// </summary>
     public interface IAuthenticationFactory : IHyakAuthenticationFactory
     {
+
+        /// <summary>
+        /// Returns IAccessToken if authentication succeeds or throws an exception if authentication fails.
+        /// </summary>
+        /// <param name="account">The azure account object</param>
+        /// <param name="environment">The azure environment object</param>
+        /// <param name="tenant">The AD tenant in most cases should be 'common'</param>
+        /// <param name="password">The AD account password</param>
+        /// <param name="promptBehavior">The prompt behavior</param>
+        /// <param name="promptAction">The prompt action used in DeviceFlow authentication</param>
+        /// <param name="optionalParameters">The optional parameters, may include TokenCache, ResourceId and CmdletContext</param>
+        /// <returns></returns>
+        IAccessToken Authenticate(
+            IAzureAccount account,
+            IAzureEnvironment environment,
+            string tenant,
+            SecureString password,
+            string promptBehavior,
+            Action<string> promptAction,
+            IDictionary<string, object> optionalParameters);
+
         /// <summary>
         /// Returns IAccessToken if authentication succeeds or throws an exception if authentication fails.
         /// </summary>
@@ -44,7 +66,6 @@ IAccessToken Authenticate(
             string promptBehavior,
             Action<string> promptAction,
             IAzureTokenCache tokenCache,
-            ICmdletContext cmdletContext,
             string resourceId = AzureEnvironment.Endpoint.ActiveDirectoryServiceEndpointResourceId);
 
         /// <summary>
@@ -65,9 +86,15 @@ IAccessToken Authenticate(
             SecureString password,
             string promptBehavior,
             Action<string> promptAction,
-            ICmdletContext cmdletContext,
             string resourceId = AzureEnvironment.Endpoint.ActiveDirectoryServiceEndpointResourceId);
 
+        /// <summary>
+        /// Get AutoRest credentials for the given context
+        /// </summary>
+        /// <param name="context">The target azure context</param>
+        /// <returns>AutoRest client credentials targeting the given context</returns>
+        ServiceClientCredentials GetServiceClientCredentials(IAzureContext context);
+
         /// <summary>
         /// Get AutoRest credentials for the given context
         /// </summary>
@@ -80,9 +107,17 @@ IAccessToken Authenticate(
         /// Get AutoRest credebntials using the given context and named endpoint
         /// </summary>
         /// <param name="context">The context to use for authentication</param>
-        /// <param name="cmdletContext">The caller cmdlet context</param>
         /// <param name="targetEndpoint">The named endpoint the AutoRest client will target</param>
         /// <returns>AutoRest client crentials targeting the given context and endpoint</returns>
+        ServiceClientCredentials GetServiceClientCredentials(IAzureContext context, string targetEndpoint);
+
+        /// <summary>
+        /// Get AutoRest credebntials using the given context and named endpoint
+        /// </summary>
+        /// <param name="context">The context to use for authentication</param>
+        /// <param name="targetEndpoint">The named endpoint the AutoRest client will target</param>
+        /// <param name="cmdletContext">The caller cmdlet context</param>
+        /// <returns>AutoRest client crentials targeting the given context and endpoint</returns>
         ServiceClientCredentials GetServiceClientCredentials(IAzureContext context, string targetEndpoint, ICmdletContext cmdletContext);
 
         /// <summary>

diff --git a/src/Authentication.Abstractions/Interfaces/IAzureTelemetry.cs b/src/Authentication.Abstractions/Interfaces/IAzureTelemetry.cs
@@ -1,32 +1,61 @@
-using System.Collections.Concurrent;
+using Newtonsoft.Json.Linq;
+using System.Threading;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
 
 namespace Microsoft.Azure.Commands.Common.Authentication.Abstractions.Interfaces
 {
     public abstract class IAzureTelemetry <T>
     {
         private ConcurrentDictionary<string, IList<T>> telemetryDataAccquirer = new ConcurrentDictionary<string, IList<T>>();
-
+
+        protected int historyKeyCount = 0;
+        public int KeysAllCount { get => historyKeyCount; }
+
+        protected int currentKeyCount = 0;
+        public int KeysCurrentCount { get => currentKeyCount; }
+
+        protected int nullCmdletContextCount = 0;
+        public int EmptyCmdletContextCount { get => nullCmdletContextCount; }
+
+        protected int keyNotFoundCount = 0;
+        public int KeyNotFoundCount { get => keyNotFoundCount; }
+
         public bool PushTelemetryRecord(ICmdletContext cmdletContext, T record)
         {
             if (cmdletContext != null && cmdletContext.IsValid && record != null)
             {
                 if (!telemetryDataAccquirer.ContainsKey(cmdletContext.CmdletId))
                 {
                     telemetryDataAccquirer[cmdletContext.CmdletId] = new List<T>();
+                    Interlocked.Increment(ref historyKeyCount);
+                    Interlocked.Increment(ref currentKeyCount);
                 }
                 telemetryDataAccquirer[cmdletContext.CmdletId].Add(record);
                 return true;
             }
+            Interlocked.Increment(ref nullCmdletContextCount);
             return false;
         }
 
         public IList<T> PopTelemetryRecord(ICmdletContext cmdletContext)
         {
-            if (cmdletContext != null && cmdletContext.IsValid && telemetryDataAccquirer.ContainsKey(cmdletContext.CmdletId))
+            if (cmdletContext != null && cmdletContext.IsValid)
+            {
+                if (telemetryDataAccquirer.ContainsKey(cmdletContext.CmdletId))
+                {
+                    telemetryDataAccquirer.TryRemove(cmdletContext.CmdletId, out IList<T> records);
+                    Interlocked.Decrement(ref currentKeyCount);
+                    return records;
+                }
+                else
+                {
+                    Interlocked.Increment(ref keyNotFoundCount);
+                }
+            }
+            else
             {
-                telemetryDataAccquirer.TryRemove(cmdletContext.CmdletId, out IList<T> records);
-                return records;
+                Interlocked.Increment(ref nullCmdletContextCount);
             }
             return null;
         }

diff --git a/src/Common/AzurePSCmdlet.cs b/src/Common/AzurePSCmdlet.cs
@@ -852,6 +852,7 @@ protected void LogQosEvent()
                 throw new NullReferenceException("AuthenticationTelemetry not registered");
             }
             _qosEvent.AuthTelemetry = authenticationTelemetry.GetTelemetryRecord(_cmdletContext);
+            WriteDebugWithTimestamp($"TotalKeyCount={authenticationTelemetry.KeysAllCount}; CurrentKeyCount={authenticationTelemetry.KeysCurrentCount}; EmptyCmdletContextCount={authenticationTelemetry.EmptyCmdletContextCount}; KeysNotFoundCount={authenticationTelemetry.KeyNotFoundCount}");
 
             if (!IsUsageMetricEnabled && (!IsErrorMetricEnabled || _qosEvent.IsSuccess))
             {

diff --git a/src/Graph.Rbac/Version1_6.20190326/ActiveDirectory/ActiveDirectoryBaseCmdlet.cs b/src/Graph.Rbac/Version1_6.20190326/ActiveDirectory/ActiveDirectoryBaseCmdlet.cs
@@ -31,7 +31,7 @@ public ActiveDirectoryClient ActiveDirectoryClient
             {
                 if (activeDirectoryClient == null)
                 {
-                    activeDirectoryClient = new ActiveDirectoryClient(DefaultProfile.DefaultContext);
+                    activeDirectoryClient = new ActiveDirectoryClient(DefaultProfile.DefaultContext, _cmdletContext);
                 }
 
                 return activeDirectoryClient;

diff --git a/src/Graph.Rbac/Version1_6.20190326/ActiveDirectory/ActiveDirectoryClient.cs b/src/Graph.Rbac/Version1_6.20190326/ActiveDirectory/ActiveDirectoryClient.cs
@@ -33,18 +33,29 @@ namespace Microsoft.Azure.Graph.RBAC.Version1_6_20190326.ActiveDirectory
 {
     public class ActiveDirectoryClient
     {
-        public ICmdletContext CmdletContext { get; set; }
-
         public GraphRbacManagementClient GraphClient { get; private set; }
 
+        /// <summary>
+        /// Creates new ActiveDirectoryClient using WindowsAzureSubscription.
+        /// </summary>
+        /// <param name="context"></param>
+        /// <param name="cmdletContext">The cmdlet context</param>
+        public ActiveDirectoryClient(IAzureContext context, ICmdletContext cmdletContext)
+        {
+            GraphClient = AzureSession.Instance.ClientFactory.CreateArmClient<GraphRbacManagementClient>(
+                context, AzureEnvironment.Endpoint.Graph, cmdletContext);
+
+            GraphClient.TenantID = context.Tenant.Id.ToString();
+        }
+
         /// <summary>
         /// Creates new ActiveDirectoryClient using WindowsAzureSubscription.
         /// </summary>
         /// <param name="context"></param>
         public ActiveDirectoryClient(IAzureContext context)
         {
             GraphClient = AzureSession.Instance.ClientFactory.CreateArmClient<GraphRbacManagementClient>(
-                context, AzureEnvironment.Endpoint.Graph, CmdletContext);
+                context, AzureEnvironment.Endpoint.Graph);
 
             GraphClient.TenantID = context.Tenant.Id.ToString();
         }

diff --git a/src/Graph.Rbac/Version1_6/ActiveDirectory/ActiveDirectoryBaseCmdlet.cs b/src/Graph.Rbac/Version1_6/ActiveDirectory/ActiveDirectoryBaseCmdlet.cs
@@ -31,7 +31,7 @@ public ActiveDirectoryClient ActiveDirectoryClient
             {
                 if (activeDirectoryClient == null)
                 {
-                    activeDirectoryClient = new ActiveDirectoryClient(DefaultProfile.DefaultContext);
+                    activeDirectoryClient = new ActiveDirectoryClient(DefaultProfile.DefaultContext, _cmdletContext);
                 }
 
                 return activeDirectoryClient;

diff --git a/src/Graph.Rbac/Version1_6/ActiveDirectory/ActiveDirectoryClient.cs b/src/Graph.Rbac/Version1_6/ActiveDirectory/ActiveDirectoryClient.cs
@@ -33,18 +33,17 @@ namespace Microsoft.Azure.Graph.RBAC.Version1_6.ActiveDirectory
 {
     public class ActiveDirectoryClient
     {
-        public ICmdletContext CmdletContext { get; set; }
-
         public GraphRbacManagementClient GraphClient { get; private set; }
 
         /// <summary>
         /// Creates new ActiveDirectoryClient using WindowsAzureSubscription.
         /// </summary>
         /// <param name="context"></param>
-        public ActiveDirectoryClient(IAzureContext context)
+        /// <param name="cmdletContext">The cmdlet context</param>
+        public ActiveDirectoryClient(IAzureContext context, ICmdletContext cmdletContext)
         {
             GraphClient = AzureSession.Instance.ClientFactory.CreateArmClient<GraphRbacManagementClient>(
-                context, AzureEnvironment.Endpoint.Graph, CmdletContext);
+                context, AzureEnvironment.Endpoint.Graph, cmdletContext);
 
             GraphClient.TenantID = context.Tenant.Id.ToString();
         }

diff --git a/src/Network/Common/NetworkClient.cs b/src/Network/Common/NetworkClient.cs
@@ -43,7 +43,11 @@ public partial class NetworkClient
 
         public Action<string> WarningLogger { get; set; }
 
-        public NetworkClient(IAzureContext context): this(AzureSession.Instance.ClientFactory.CreateArmClient<NetworkManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager, AzureCmdletContext.CmdletNone))
+        public NetworkClient(IAzureContext context): this(AzureSession.Instance.ClientFactory.CreateArmClient<NetworkManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager))
+        {
+        }
+
+        public NetworkClient(IAzureContext context, ICmdletContext cmdletContext) : this(AzureSession.Instance.ClientFactory.CreateArmClient<NetworkManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager, cmdletContext))
         {
         }
 

diff --git a/src/ResourceManager/Version2016_09_01/ArgumentCompleters/LocationCompleter.cs b/src/ResourceManager/Version2016_09_01/ArgumentCompleters/LocationCompleter.cs
@@ -50,7 +50,7 @@ protected static IDictionary<string, ICollection<string>> ResourceTypeLocationDi
                     {
                         try
                         {
-                            IResourceManagementClient client = AzureSession.Instance.ClientFactory.CreateArmClient<ResourceManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager, AzureCmdletContext.CmdletNone);
+                            IResourceManagementClient client = AzureSession.Instance.ClientFactory.CreateArmClient<ResourceManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager);
                             var allProviders = client.Providers.ListAsync();
                             if (_timeout == -1)
                             {

diff --git a/src/ResourceManager/Version2016_09_01/ArgumentCompleters/ResourceGroupCompleter.cs b/src/ResourceManager/Version2016_09_01/ArgumentCompleters/ResourceGroupCompleter.cs
@@ -48,7 +48,7 @@ protected static IList<String> ResourceGroupNames
                         var tempResourceGroupList = new List<string>();
                         try
                         {
-                            var client = AzureSession.Instance.ClientFactory.CreateArmClient<ResourceManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager, AzureCmdletContext.CmdletNone);
+                            var client = AzureSession.Instance.ClientFactory.CreateArmClient<ResourceManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager);
                             // Retrieve only the first page of ResourceGroups to display to the user
                             var resourceGroups = client.ResourceGroups.ListAsync();
 

diff --git a/src/ResourceManager/Version2016_09_01/ArgumentCompleters/ResourceIdCompleter.cs b/src/ResourceManager/Version2016_09_01/ArgumentCompleters/ResourceIdCompleter.cs
@@ -59,7 +59,7 @@ public static IEnumerable<string> GetResourceIds(string resourceType)
                     return Cache[contextHash].ResourceInfoList;
                 }
 
-                var client = AzureSession.Instance.ClientFactory.CreateArmClient<ResourceManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager, AzureCmdletContext.CmdletNone);
+                var client = AzureSession.Instance.ClientFactory.CreateArmClient<ResourceManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager);
 
                 var odata = new ODataQuery<GenericResourceFilter>(r => r.ResourceType == resourceType);
 

diff --git a/src/ResourceManager/Version2016_09_01/ArgumentCompleters/ResourceNameCompleter.cs b/src/ResourceManager/Version2016_09_01/ArgumentCompleters/ResourceNameCompleter.cs
@@ -121,7 +121,7 @@ public static string CreateFilter(
         public static List<ResourceIdentifier> GetResourceIdsFromClient(string resourceType, string resourceGroupName)
         {
             IAzureContext context = AzureRmProfileProvider.Instance?.Profile?.DefaultContext;
-            IResourceManagementClient client = AzureSession.Instance.ClientFactory.CreateArmClient<ResourceManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager, AzureCmdletContext.CmdletNone);
+            IResourceManagementClient client = AzureSession.Instance.ClientFactory.CreateArmClient<ResourceManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager);
             var odataQuery = new ODataQuery<GenericResourceFilter>(r => r.ResourceType == resourceType);
 
             var allProviders = string.IsNullOrWhiteSpace(resourceGroupName)

diff --git a/src/ResourceManager/Version2016_09_01/ArgumentCompleters/ResourceTypeCompleter.cs b/src/ResourceManager/Version2016_09_01/ArgumentCompleters/ResourceTypeCompleter.cs
@@ -48,7 +48,7 @@ protected static IList<String> ResourceTypes
                         var tempResourceTypeList = new List<string>();
                         try
                         {
-                            var client = AzureSession.Instance.ClientFactory.CreateArmClient<ResourceManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager, AzureCmdletContext.CmdletNone);
+                            var client = AzureSession.Instance.ClientFactory.CreateArmClient<ResourceManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager);
                             var resourceTypes = new List<Provider>();
                             var task = client.Providers.ListAsync();
                             if (task.Wait(TimeSpan.FromSeconds(_timeout)))

diff --git a/src/ResourceManager/Version2016_09_01/ArgumentCompleters/ScopeCompleter.cs b/src/ResourceManager/Version2016_09_01/ArgumentCompleters/ScopeCompleter.cs
@@ -45,7 +45,7 @@ protected static IList<String> Scopes
                         var tempScopeList = new List<string>();
                         try
                         {
-                            var client = AzureSession.Instance.ClientFactory.CreateArmClient<ResourceManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager, AzureCmdletContext.CmdletNone);
+                            var client = AzureSession.Instance.ClientFactory.CreateArmClient<ResourceManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager);
                             // Retrieve only the first page of ResourceGroups to use for scopes
                             var resourceGroups = client.ResourceGroups.ListAsync();
 

diff --git a/src/ResourceManager/Version2016_09_01/Tags/TagsClient.cs b/src/ResourceManager/Version2016_09_01/Tags/TagsClient.cs
@@ -38,7 +38,16 @@ public class TagsClient
         /// Creates new tags client instance.
         /// </summary>
         /// <param name="context">The Azure context instance</param>
-        public TagsClient(IAzureContext context): this(AzureSession.Instance.ClientFactory.CreateArmClient<ResourceManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager, AzureCmdletContext.CmdletNone))
+        public TagsClient(IAzureContext context): this(AzureSession.Instance.ClientFactory.CreateArmClient<ResourceManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager))
+        {
+        }
+
+        /// <summary>
+        /// Creates new tags client instance.
+        /// </summary>
+        /// <param name="context">The Azure context instance</param>
+        /// <param name="cmdletContext">The cmdlet context</param>
+        public TagsClient(IAzureContext context, ICmdletContext cmdletContext) : this(AzureSession.Instance.ClientFactory.CreateArmClient<ResourceManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager, cmdletContext))
         {
         }
 

diff --git a/src/ResourceManager/Version2016_09_01/Utilities/SubscriptionAndTenantHelper.cs b/src/ResourceManager/Version2016_09_01/Utilities/SubscriptionAndTenantHelper.cs
@@ -14,6 +14,7 @@
 
 using Microsoft.Azure.Commands.Common.Authentication;
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
+using Microsoft.Azure.Commands.Common.Authentication.Abstractions.Extensions;
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions.Interfaces;
 using Microsoft.Azure.Internal.Subscriptions;
 using Microsoft.Azure.Internal.Subscriptions.Models;
@@ -36,7 +37,7 @@ internal static IAccessToken AcquireAccessToken(IAzureAccount account, IAzureEnv
                null,
                ShowDialog.Never,
                null,
-               cmdletContext);
+               cmdletContext.ToExtensibleParameters());
         }
 
         internal static Dictionary<string, AzureSubscription> GetTenantsForSubscriptions(List<string> subscriptionIds, IAzureContext defaultContext, ICmdletContext cmdletContext)