Sync tools folder from main branch to generation branch (#21053)

Co-authored-by: azurepowershell <[email protected]>

Author: azure-powershell-bot

src/Az.Shared.props

@@ -8,9 +8,20 @@
 
   <!-- Build -->
   <PropertyGroup>
-    <GenerateAssemblyInfo>false</GenerateAssemblyInfo>
+    <GenerateAssemblyCompanyAttribute>false</GenerateAssemblyCompanyAttribute>
+    <GenerateAssemblyConfigurationAttribute>false</GenerateAssemblyConfigurationAttribute>
+    <GenerateAssemblyFileVersionAttribute>false</GenerateAssemblyFileVersionAttribute>
+    <GenerateAssemblyInformationalVersionAttribute>false</GenerateAssemblyInformationalVersionAttribute>
+    <GenerateAssemblyProductAttribute>false</GenerateAssemblyProductAttribute>
+    <GenerateAssemblyTitleAttribute>false</GenerateAssemblyTitleAttribute>
+    <GenerateAssemblyVersionAttribute>false</GenerateAssemblyVersionAttribute>
     <TreatWarningsAsErrors>true</TreatWarningsAsErrors>
     <WarningsAsErrors />
   </PropertyGroup>
 
-</Project>
+  <!-- include shared source code -->
+  <ItemGroup>
+    <Compile Include="$(MSBuildThisFileDirectory)shared\*.cs" />
+  </ItemGroup>
+
+</Project>

src/Az.props

@@ -24,7 +24,7 @@
     <SignAssembly>true</SignAssembly>
     <DelaySign>true</DelaySign>
     <AssemblyOriginatorKeyFile>$(RepoSrc)MSSharedLibKey.snk</AssemblyOriginatorKeyFile>
-    <DefineConstants>TRACE;RELEASE;NETSTANDARD;SIGN</DefineConstants>
+    <DefineConstants>TRACE;RELEASE;NETSTANDARD;SIGN;$(TestCoverage)</DefineConstants>
   </PropertyGroup>
 
   <!-- Resources -->
@@ -34,8 +34,6 @@
       <AutoGen>true</AutoGen>
       <DependentUpon>Resources.resx</DependentUpon>
     </Compile>
-    <!-- include shared source code -->
-    <Compile Include="$(MSBuildThisFileDirectory)shared\*.cs" />
   </ItemGroup>
 
   <ItemGroup>
@@ -53,4 +51,4 @@
     <None Update="*.ps1xml" CopyToOutputDirectory="PreserveNewest" />
   </ItemGroup>
 
-</Project>
+</Project>

src/shared/ConfigKeys.cs

@@ -31,5 +31,6 @@ internal static class ConfigKeys
         public const string EnableTestCoverage = "EnableTestCoverage";
         public const string EnableLoginByWam = "EnableLoginByWam";
         public const string TestCoverageLocation = "TestCoverageLocation";
+        public const string DisplayRegionIdentified = "DisplayRegionIdentified";
     }
 }

tools/Az.Tools.Predictor/Az.Tools.Predictor.Test/AzPredictorServiceTests.cs

@@ -229,7 +229,7 @@ public void VerifyNoPrediction(string userInput)
             Assert.Equal(0, actual.Count);
 
             actual = this._noPredictorService.GetSuggestion(predictionContext, 1, 1, CancellationToken.None);
-            Assert.Null(actual);
+            Assert.Empty(actual.PredictiveSuggestions);
         }
 
         /// <summary>

tools/Az.Tools.Predictor/Az.Tools.Predictor.Test/AzPredictorTelemetryTests.cs

@@ -59,17 +59,20 @@ public void VerifyUnescapingText(string expected, string input)
         [InlineData("Get-AzResourceGroup -Location:*** -Name ***", "Get-AzResourceGroup -Name Test -Location:WestUS2")]
         // skip the positional parameter
         [InlineData("Get-AzResourceGroup -Location:***", "Get-AzResourceGroup Test -Location:WestUS2")]
+        [InlineData("Get-AzResourceGroup", "Get-AzResourceGroup Test WestUS2")]
+        // Take the last command in a pipe
+        [InlineData("Set-AzContext", "$context | Set-AzContext")]
+        // Take the command at the right of assignment
+        [InlineData("Get-AzRoleAssignment -ObjectId ***", "$a=Get-AzRoleAssignment -ObjectId xxx")]
+        // Take the command inside parentheses
+        [InlineData("Get-AzRoleAssignment -ObjectId ***", "Remove-AzRoleAssignment -RoleDefinitionId (Get-AzRoleAssignment -ObjectId xxx")]
         [InlineData("git", "git status")]
         [InlineData("customCommand -Parameter ***", "customCommand -Parameter Value")]
         [InlineData("customCommand", "customCommand /Parameter Value")]
         [InlineData("customCommand", "customCommand --Parameter Value")]
         public void VerifyMaskCommandLine(string expected, string input)
         {
-            var asts = Parser.ParseInput(input, out _, out _);
-            var allNestedAsts = asts?.FindAll((ast) => ast is CommandAst, true);
-            var commandAst = allNestedAsts?.LastOrDefault() as CommandAst;
-
-            string maskedCommandLine = CommandLineUtilities.MaskCommandLine(commandAst);
+            string maskedCommandLine = CommandLineUtilities.MaskCommandLine(input);
 
             Assert.Equal(expected, maskedCommandLine);
         }
@@ -82,7 +85,7 @@ public void VerifyMaskCommandLine(string expected, string input)
         [InlineData("Set-AzContext -Subscription 'xxxx-xxxx-xxxx-xxxx' -Tenant {String}", null)]
         [InlineData("Set-AzContext -Subscription $subscription -Tenant [[String]]", null)]
         [InlineData("Set-AzContext -Subscription 'xxxx-xxxx-xxxx-xxxx' -Tenant TenantName", null)]
-        [InlineData("Get-AzContext | Set-AzContext", "Get-AzContext")]
+        [InlineData("Get-AzContext | Set-AzContext", "Set-AzContext")]
         public void VerifyGetCommandAst(string input, string expected)
         {
             if (expected == null)

tools/Az.Tools.Predictor/Az.Tools.Predictor.Test/CommandLineTests.cs

@@ -12,6 +12,7 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
+using System;
 using System.Collections.Generic;
 using System.Management.Automation.Language;
 using System.Management.Automation.Subsystem.Prediction;
@@ -71,8 +72,8 @@ public MockAzPredictorService(string history, IList<PredictiveCommand> suggestio
             }
         }
 
-        /// <inheritdoc/>
-        public override bool IsSupportedCommand(string cmd) => IsRecognizedCommand(cmd);
+        // <inheritdoc/>
+        // public override bool IsSupportedCommand(string cmd) => IsRecognizedCommand(cmd) || cmd.IndexOf(AzPredictorConstants.AzCommandMoniker, StringComparison.OrdinalIgnoreCase) > 0;
 
         /// <inheritdoc/>
         public override Task<bool?> RequestPredictionsAsync(IEnumerable<string> commands, string requestId, CancellationToken cancellationToken)

tools/Az.Tools.Predictor/Az.Tools.Predictor.Test/Mocks/MockAzPredictorService.cs

@@ -127,7 +127,7 @@ public Version ModuleVersion
         }
 
         /// <inheritdoc/>
-        public bool IsInternal { get; internal set; }
+        public bool IsInternal { get; private set; }
 
         public AzContext(IPowerShellRuntime powerShellRuntime) => _powerShellRuntime
              = powerShellRuntime;
@@ -138,11 +138,7 @@ public void UpdateContext()
             AzVersion = GetAzVersion();
             RawUserId = GetUserAccountId();
             HashUserId = GenerateSha256HashString(RawUserId);
-
-            if (!IsInternal)
-            {
-                IsInternal = RawUserId.EndsWith(AzContext.InternalUserSuffix, StringComparison.OrdinalIgnoreCase);
-            }
+            IsInternal = RawUserId.EndsWith(AzContext.InternalUserSuffix, StringComparison.OrdinalIgnoreCase);
         }
 
         internal string RawUserId { get; set; }

tools/Az.Tools.Predictor/Az.Tools.Predictor/AzContext.cs

@@ -93,10 +93,7 @@ public AzPredictor()
             Task.Run(() =>
                     {
                         _settings = Settings.GetSettings();
-                        _azContext = new AzContext(_powerShellRuntime)
-                        {
-                            IsInternal = (_settings.SetAsInternal == true) ? true : false,
-                        };
+                        _azContext = new AzContext(_powerShellRuntime);
 
                         _azContext.UpdateContext();
                         // This will run the script in the right context.
@@ -279,7 +276,13 @@ public void OnCommandLineExecuted(PredictionClient client, string commandLine, b
                     parsedResult = GetAstAndMaskedCommandLine(commandLine);
                 }
 
-                _telemetryClient.OnHistory(new HistoryTelemetryData(client, parsedResult.MaskedCommandLine ?? AzPredictorConstants.CommandPlaceholder, success));
+                // We only collect the masked command line when the command is supported.
+                // This is to avoid the case when the user copies and pasted some sensitive data into the command line.
+
+                var collectedCommandLine = parsedResult.IsSupported ?
+                    parsedResult.MaskedCommandLine ?? AzPredictorConstants.CommandPlaceholder :
+                    AzPredictorConstants.CommandPlaceholder;
+                _telemetryClient.OnHistory(new HistoryTelemetryData(client, collectedCommandLine, success));
                 _commandLineExecutedCompletion?.SetResult();
             });
 
@@ -315,7 +318,10 @@ public SuggestionPackage GetSuggestion(PredictionClient client, PredictionContex
                 }
                 finally
                 {
-                    _telemetryClient.OnGetSuggestion(new GetSuggestionTelemetryData(client, localSuggestionSessionId, context.InputAst,
+                    _telemetryClient.OnGetSuggestion(new GetSuggestionTelemetryData(client,
+                            localSuggestionSessionId,
+                            suggestions?.CommandAst,
+                            _service.IsSupportedCommand(suggestions?.CommandAst?.GetCommandName()),
                             suggestions,
                             cancellationToken.IsCancellationRequested,
                             exception));

tools/Az.Tools.Predictor/Az.Tools.Predictor/AzPredictor.cs

@@ -160,16 +160,21 @@ public virtual CommandLineSuggestion GetSuggestion(PredictionContext context, in
                 }
             }
 
+            CommandLineSuggestion earlyReturnResult = new()
+            {
+                CommandAst = commandAst
+            };
+
             if (commandAst == null)
             {
-                return null;
+                return earlyReturnResult;
             }
 
             var commandName = commandAst.GetCommandName();
 
             if (string.IsNullOrWhiteSpace(commandName))
             {
-                return null;
+                return earlyReturnResult;
             }
 
             ParameterSet inputParameterSet = null;
@@ -182,7 +187,7 @@ public virtual CommandLineSuggestion GetSuggestion(PredictionContext context, in
             {
                 // We only ignore the exception when the command name is not supported.
                 // We want to collect the telemetry about the exception how common it is for the format we don't support.
-                return null;
+                return earlyReturnResult;
             }
 
             cancellationToken.ThrowIfCancellationRequested();
@@ -250,6 +255,13 @@ public virtual CommandLineSuggestion GetSuggestion(PredictionContext context, in
                 }
             }
 
+            if (result is null)
+            {
+                result = new CommandLineSuggestion();
+            }
+
+            result.CommandAst = commandAst;
+
             return result;
         }
 
@@ -324,8 +336,10 @@ public virtual void RecordHistory(CommandAst history) =>
 
         /// <inheritdoc/>
         public virtual bool IsSupportedCommand(string cmd) => IsRecognizedCommand(cmd)
-            && !_surveyCmdlets.Any(cmdlet => cmdlet.Command.StartsWith(cmd, StringComparison.OrdinalIgnoreCase)) // the survey cmdlets aren't in the normal az command flow, so mark them as unsupported.
-            && cmd.IndexOf(AzPredictorConstants.AzCommandMoniker) > 0; // This is the Az cmdlet.
+            || (!string.IsNullOrWhiteSpace(cmd)
+                && !_surveyCmdlets.Any(cmdlet => cmdlet.Command.StartsWith(cmd, StringComparison.OrdinalIgnoreCase)) // the survey cmdlets aren't in the normal az command flow, so mark them as unsupported.
+                && cmd.IndexOf(AzPredictorConstants.AzCommandMoniker, StringComparison.OrdinalIgnoreCase) > 0 // This is the Az cmdlet.
+            );
 
         /// <summary>
         /// Checks whether the given <paramref name="cmd" /> is in the command list.

tools/Az.Tools.Predictor/Az.Tools.Predictor/AzPredictorService.cs

@@ -1,6 +1,5 @@
 {
   "maxAllowedCommandDuplicate": 1,
   "serviceUri": "https://app.aladdin.microsoft.com/api/v2.0",
-  "setAsInternal": false,
   "suggestionCount": 7
 }

tools/Az.Tools.Predictor/Az.Tools.Predictor/AzPredictorSettings.json

@@ -15,6 +15,7 @@
 using Microsoft.Azure.PowerShell.Tools.AzPredictor.Utilities;
 using System;
 using System.Collections.Generic;
+using System.Management.Automation.Language;
 using System.Management.Automation.Subsystem.Prediction;
 
 namespace Microsoft.Azure.PowerShell.Tools.AzPredictor
@@ -25,7 +26,7 @@ namespace Microsoft.Azure.PowerShell.Tools.AzPredictor
     /// <remarks>
     /// Because the performance requirement in <see cref="AzPredictor.GetSuggestion"/>,
     /// it contains lists of each piece of information, for example, a collection of predictive suggestion and a list of
-    /// suggestion sources. Note that the count of each list should be the same. And each element in the list corresonds to
+    /// suggestion sources. Note that the count of each list should be the same. And each element in the list corresponds to
     /// the element in other list at the same index.
     /// </remarks>
     public sealed class CommandLineSuggestion
@@ -38,6 +39,11 @@ public sealed class CommandLineSuggestion
 
         private readonly List<PredictiveSuggestion> _predictiveSuggestions = new List<PredictiveSuggestion>(CommandLineSuggestion.CollectionDefaultCapacity);
 
+        /// <summary>
+        /// Gets or sets the AST the suggestions are provided for. Note this is not always the same as the whole user input.
+        /// </summary>
+        public CommandAst CommandAst { get; set; }
+
         /// <summary>
         /// Gets the suggestions returned to show to the user. This can be adjusted from <see cref="SourceTexts"/> based on
         /// the user input.

tools/Az.Tools.Predictor/Az.Tools.Predictor/CommandLineSuggestion.cs

@@ -37,10 +37,7 @@ protected override void EndProcessing()
             TelemetryClient telemetryClient = TelemetryUtilities.CreateApplicationInsightTelemetryClient();
             var settings = Settings.GetSettings();
             var nestedPowerShellRuntime = new PowerShellRuntime();
-            var azContext = new AzContext(nestedPowerShellRuntime)
-            {
-                IsInternal = (settings.SetAsInternal == true) ? true : false,
-            };
+            var azContext = new AzContext(nestedPowerShellRuntime);
 
             var invocation = MyInvocation;
             var command = invocation.MyCommand;

tools/Az.Tools.Predictor/Az.Tools.Predictor/Commands/BasePSCmdlet.cs

@@ -45,7 +45,7 @@ protected override void ProcessRecord()
             int minIdNumber = 0;
             int maxIdNumber = 1000000;
             // Format the integer into a 6-digit string, adding 0 to the left if needed, e.g. 123 -> "000123", 123456 -> "123456".
-            // See https://docs.microsoft.com/en-us/dotnet/standard/base-types/standard-numeric-format-strings
+            // See https://learn.microsoft.com/en-us/dotnet/standard/base-types/standard-numeric-format-strings
             var surveyId = random.Next(minIdNumber, maxIdNumber).ToString("D6", CultureInfo.InvariantCulture);
 
             var link = string.Format(OpenAzPredictorSurvey._SurveyLinkFormat, surveyId, CultureInfo.InvariantCulture);

tools/Az.Tools.Predictor/Az.Tools.Predictor/Commands/OpenAzPredictorSurvey.cs

@@ -32,7 +32,7 @@ public interface IAzPredictorService
         /// <param name="suggestionCount">The number of suggestion to return.</param>
         /// <param name="cancellationToken">The cancellation token.</param>
         /// <param name="maxAllowedCommandDuplicate">The maximum amount of the same commnds in the list of predictions.</param>
-        /// <returns>The suggestions for <paramref name="context"/>. The maximum number of suggestions is <paramref name="suggestionCount"/>. A null will be returned if there the user input context isn't valid/supported at all.</returns>
+        /// <returns>The suggestions for <paramref name="context"/>. The maximum number of suggestions is <paramref name="suggestionCount"/>.</returns>
         public CommandLineSuggestion GetSuggestion(PredictionContext context, int suggestionCount, int maxAllowedCommandDuplicate, CancellationToken cancellationToken);
 
         /// <summary>
@@ -51,8 +51,8 @@ public interface IAzPredictorService
         public void RecordHistory(CommandAst history);
 
         /// <summary>
-        /// Return true if command is part of known set of Az cmdlets, false otherwise.
+        /// Return true if command of the name is part of known set of Az cmdlets, false otherwise.
         /// </summary>
-        public bool IsSupportedCommand(string cmd);
+        public bool IsSupportedCommand(string commandName);
     }
 }

tools/Az.Tools.Predictor/Az.Tools.Predictor/IAzPredictorService.cs

@@ -303,14 +303,14 @@ protected void SendAggregatedTelemetryData()
 
             if (suggestions.Count > 0)
             {
-                properties.Add("Suggestion", JsonSerializer.Serialize(suggestions, JsonUtilities.TelemetrySerializerOptions));
+                properties.Add(GetSuggestionTelemetryData.PropertyNamePrediction, JsonSerializer.Serialize(suggestions, JsonUtilities.TelemetrySerializerOptions));
             }
 
             if (aggregatedData.CommandLine != null)
             {
                 // Add the command history data.
                 properties.Add(HistoryTelemetryData.PropertyNameSuccess, aggregatedData.IsCommandSuccess.ToString(CultureInfo.InvariantCulture));
-                properties.Add(HistoryTelemetryData.PropertyNameHistory, aggregatedData.CommandLine);
+                properties.Add(HistoryTelemetryData.PropertyNameCommand, aggregatedData.CommandLine);
             }
 
             SendTelemetry($"{TelemetryUtilities.TelemetryEventPrefix}/Aggregation", properties);
@@ -321,6 +321,12 @@ protected void SendAggregatedTelemetryData()
         /// </summary>
         protected virtual void SendTelemetry(string eventName, IDictionary<string, string> properties)
         {
+            if (this._azContext.IsInternal)
+            {
+                // To better filter internal telemetry, we append "_INTERNAL" to the event name.
+                eventName = eventName + "_INTERNAL";
+            }
+
             _telemetryClient.TrackEvent(eventName, properties);
         }
 
@@ -442,7 +448,9 @@ private void ProcessTelemetryData(GetSuggestionTelemetryData telemetryData)
                 }
             }
 
-            var maskedUserInput = CommandLineUtilities.MaskCommandLine(telemetryData.UserInput.FindAll((ast) => ast is CommandAst, true).LastOrDefault() as CommandAst);
+            var maskedUserInput = telemetryData.IsSupported ?
+                CommandLineUtilities.MaskCommandLine(CommandLineUtilities.GetCommandAst(telemetryData.UserInput)) :
+                AzPredictorConstants.CommandPlaceholder;
 
             var suggestionSession = new SuggestionSession()
             {
@@ -500,7 +508,7 @@ private void ProcessTelemetryData(SuggestionDisplayedTelemetryData telemetryData
             if (CachedAggregatedTelemetryData.EstimateSuggestionSessionSize >= AzPredictorTelemetryClient.MaxPropertyValueSizeWithBuffer)
             {
                 suggestionSession = SendAggregateTelemetryDataDuringSuggestionCycle(telemetryData, telemetryData.SuggestionSessionId);
-                suggestionSession.IsSuggestionComplete = false; // This continue from the previous suggestionsession. So mark it as incomplete.
+                suggestionSession.IsSuggestionComplete = false; // This continue from the previous suggestion session. So mark it as incomplete.
             }
 
             suggestionSession.DisplayMode = telemetryData.DisplayMode;