Make the default sensitive history scrubbing function a little smarter

PSReadLine/History.cs

@@ -12,6 +12,7 @@
 using System.Text;
 using System.Text.RegularExpressions;
 using System.Threading;
+using System.Management.Automation.Language;
 using Microsoft.PowerShell.PSReadLine;
 
 namespace Microsoft.PowerShell
@@ -116,6 +117,20 @@ public class HistoryItem
             "password|asplaintext|token|apikey|secret",
             RegexOptions.Compiled | RegexOptions.IgnoreCase);
 
+        private static readonly HashSet<string> s_SecretMgmtCommands = new(StringComparer.OrdinalIgnoreCase)
+        {
+            "Get-Secret",
+            "Get-SecretInfo",
+            "Get-SecretVault",
+            "Register-SecretVault",
+            "Remove-Secret",
+            "Set-SecretInfo",
+            "Set-SecretVaultDefault",
+            "Test-SecretVault",
+            "Unlock-SecretVault",
+            "Unregister-SecretVault"
+        };
+
         private void ClearSavedCurrentLine()
         {
             _savedCurrentLine.CommandLine = null;
@@ -474,11 +489,169 @@ void UpdateHistoryFromFile(IEnumerable<string> historyLines, bool fromDifferentS
             }
         }
 
+        private static bool IsOnLeftSideOfAnAssignment(Ast ast, out Ast rhs)
+        {
+            bool result = false;
+            rhs = null;
+
+            do
+            {
+                if (ast.Parent is AssignmentStatementAst assignment)
+                {
+                    rhs = assignment.Right;
+                    result = ReferenceEquals(assignment.Left, ast);
+
+                    break;
+                }
+
+                ast = ast.Parent;
+            }
+            while (ast.Parent is not null);
+
+            return result;
+        }
+
+        private static bool IsSecretMgmtCommand(StringConstantExpressionAst strConst, out CommandAst command)
+        {
+            bool result = false;
+            command = strConst.Parent as CommandAst;
+
+            if (command is not null)
+            {
+                result = ReferenceEquals(command.CommandElements[0], strConst)
+                    && s_SecretMgmtCommands.Contains(strConst.Value);
+            }
+
+            return result;
+        }
+
+        private static ExpressionAst GetArgumentForParameter(CommandParameterAst param)
+        {
+            if (param.Argument is not null)
+            {
+                return param.Argument;
+            }
+
+            var command = (CommandAst)param.Parent;
+            int index = 1;
+            for (; index < command.CommandElements.Count; index++)
+            {
+                if (ReferenceEquals(command.CommandElements[index], param))
+                {
+                    break;
+                }
+            }
+
+            int argIndex = index + 1;
+            if (argIndex < command.CommandElements.Count
+                && command.CommandElements[argIndex] is ExpressionAst arg)
+            {
+                return arg;
+            }
+
+            return null;
+        }
+
         public static AddToHistoryOption GetDefaultAddToHistoryOption(string line)
         {
-            return s_sensitivePattern.IsMatch(line)
-                ? AddToHistoryOption.MemoryOnly
-                : AddToHistoryOption.MemoryAndFile;
+            if (string.IsNullOrEmpty(line))
+            {
+                return AddToHistoryOption.SkipAdding;
+            }
+
+            Match match = s_sensitivePattern.Match(line);
+            if (ReferenceEquals(match, Match.Empty))
+            {
+                return AddToHistoryOption.MemoryAndFile;
+            }
+
+            bool isSensitive = false;
+            Ast ast = string.Equals(_singleton._ast?.Extent.Text, line)
+                ? _singleton._ast
+                : Parser.ParseInput(line, out _, out _);
+
+            do
+            {
+                int start = match.Index;
+                int end = start + match.Length;
+
+                IEnumerable<Ast> asts = ast.FindAll(
+                    ast => ast.Extent.StartOffset <= start && ast.Extent.EndOffset >= end,
+                    searchNestedScriptBlocks: true);
+
+                Ast innerAst = asts.Last();
+                switch (innerAst)
+                {
+                    case VariableExpressionAst:
+                        // It's a variable with sensitive name. Using the variable is fine, but assigning to
+                        // the variable could potentially expose sensitive content.
+                        // If it appears on the left-hand-side of an assignment, and the right-hand-side is
+                        // not a command invocation, we consider it sensitive.
+                        // e.g. `$token = Get-Secret` vs. `$token = 'token-text'` or `$token, $url = ...`
+                        isSensitive = IsOnLeftSideOfAnAssignment(innerAst, out Ast rhs)
+                            && rhs is not PipelineAst;
+
+                        if (!isSensitive)
+                        {
+                            match = match.NextMatch();
+                        }
+                        break;
+
+                    case StringConstantExpressionAst strConst:
+                        // If it's not a command name, or it's not one of the secret management commands that
+                        // we can ignore, we consider it sensitive.
+                        isSensitive = !IsSecretMgmtCommand(strConst, out CommandAst command);
+
+                        if (!isSensitive)
+                        {
+                            // We can safely skip the whole command text.
+                            match = s_sensitivePattern.Match(line, command.Extent.EndOffset);
+                        }
+                        break;
+
+                    case CommandParameterAst param:
+                        // Special-case the '-AsPlainText' parameter.
+                        if (string.Equals(param.ParameterName, "AsPlainText"))
+                        {
+                            isSensitive = true;
+                            break;
+                        }
+
+                        ExpressionAst arg = GetArgumentForParameter(param);
+                        if (arg is null)
+                        {
+                            // If no argument is found following the parameter, then it could be a switching parameter
+                            // such as '-UseDefaultPassword' or '-SaveToken', which we assume will not expose sensitive information.
+                            match = match.NextMatch();
+                        }
+                        else if (arg is VariableExpressionAst)
+                        {
+                            // Argument is a variable. It's fine to use a variable for a senstive parameter.
+                            // e.g. `Invoke-WebRequest -Token $token`
+                            match = s_sensitivePattern.Match(line, arg.Extent.EndOffset);
+                        }
+                        else if (arg is ParenExpressionAst paren
+                            && paren.Pipeline is PipelineAst pipeline
+                            && pipeline.PipelineElements[0] is not CommandExpressionAst)
+                        {
+                            // Argument is a command invocation, such as `Invoke-WebRequest -Token (Get-Secret)`.
+                            match = match.NextMatch();
+                        }
+                        else
+                        {
+                            // We consider all other arguments sensitive.
+                            isSensitive = true;
+                        }
+                        break;
+
+                    default:
+                        isSensitive = true;
+                        break;
+                }
+            }
+            while (!isSensitive && !ReferenceEquals(match, Match.Empty));
+
+            return isSensitive ? AddToHistoryOption.MemoryOnly : AddToHistoryOption.MemoryAndFile;
         }
 
         /// <summary>

test/HistoryTest.cs

@@ -89,7 +89,7 @@ public void ParallelHistorySaving()
         }
 
         [SkippableFact]
-        public void SensitiveHistoryDefaultBehavior()
+        public void SensitiveHistoryDefaultBehavior_One()
         {
             TestSetup(KeyMode.Cmd);
 
@@ -163,6 +163,92 @@ public void SensitiveHistoryDefaultBehavior()
             }
         }
 
+        [SkippableFact]
+        public void SensitiveHistoryDefaultBehavior_Two()
+        {
+            TestSetup(KeyMode.Cmd);
+
+            // Clear history
+            SetHistory();
+
+            var options = PSConsoleReadLine.GetOptions();
+            var oldHistoryFilePath = options.HistorySavePath;
+            var oldHistorySaveStyle = options.HistorySaveStyle;
+
+            // AddToHistoryHandler should be set to the default handler.
+            Assert.Same(PSConsoleReadLineOptions.DefaultAddToHistoryHandler, options.AddToHistoryHandler);
+
+            var newHistoryFilePath = Path.GetTempFileName();
+            var newHistorySaveStyle = HistorySaveStyle.SaveIncrementally;
+
+            string[] expectedHistoryItems = new[] {
+                "$token = 'abcd'", // Assign expr-value to sensitive variable. Not saved to file.
+                "Set-Secret abc $mySecret", // 'Set-Secret' will not be save to file.
+                "ConvertTo-SecureString stringValue -AsPlainText", // '-AsPlainText' is an alert. Not saved to file.
+                "Get-Secret PSGalleryApiKey -AsPlainText", // For eligible secret-mgmt command, the whole command is skipped, so '-AsPlainText' here is OK.
+                "$token = Get-Secret -Name github-token -Vault MySecret",
+                "[MyType]::CallRestAPI($token, $url, $args)",
+                "$template -f $token",
+                "Invoke-RestCall $url -UseDefaultToken",
+                "Publish-Module -NuGetApiKey $apikey",
+                "Publish-Module -NuGetApiKey (Get-Secret PSGalleryApiKey -AsPlainText)",
+                "Publish-Module -NuGetApiKey (Get-NewSecret -Name apikey)", // 'Get-NewSecret' is not in our allow-list. Not saved to file.
+                "Send-HeartBeat -UseDefaultToken",
+                "Send-HeartBeat -password $pass -SavePassword",
+                "Invoke-WebRequest -Token xxx", // Expr-value as argument to '-Token'. Not saved to file.
+                "Invoke-WebRequest -Token (2+2)", // Expr-value as argument to '-Token'. Not saved to file.
+                "Get-SecretInfo -Name mytoken; Get-SecretVault; Register-SecretVault; Remove-Secret apikey",
+                "Get-SecretInfo -Name mytoken; Get-SecretVault; Register-SecretVault; Remove-Secret apikey; Set-Secret", // 'Set-Secret' Not saved to file.
+                "Set-SecretInfo -Name apikey; Set-SecretVaultDefault; Test-SecretVault; Unlock-SecretVault -password $pwd; Unregister-SecretVault -SecretVault vaultInfo",
+                "Get-ResultFromTwo -Secret1 (Get-Secret -Name blah -AsPlainText) -Secret2 $secret2",
+                "Get-ResultFromTwo -Secret1 (Get-Secret -Name blah -AsPlainText) -Secret2 sdv87ysdfayf798hfasd8f7ha" // '-Secret2' has expr-value argument. Not saved to file.
+            };
+
+            string[] expectedSavedItems = new[] {
+                "Get-Secret PSGalleryApiKey -AsPlainText",
+                "$token = Get-Secret -Name github-token -Vault MySecret",
+                "[MyType]::CallRestAPI($token, $url, $args)",
+                "$template -f $token",
+                "Invoke-RestCall $url -UseDefaultToken",
+                "Publish-Module -NuGetApiKey $apikey",
+                "Publish-Module -NuGetApiKey (Get-Secret PSGalleryApiKey -AsPlainText)",
+                "Send-HeartBeat -UseDefaultToken",
+                "Send-HeartBeat -password $pass -SavePassword",
+                "Get-SecretInfo -Name mytoken; Get-SecretVault; Register-SecretVault; Remove-Secret apikey",
+                "Set-SecretInfo -Name apikey; Set-SecretVaultDefault; Test-SecretVault; Unlock-SecretVault -password $pwd; Unregister-SecretVault -SecretVault vaultInfo",
+                "Get-ResultFromTwo -Secret1 (Get-Secret -Name blah -AsPlainText) -Secret2 $secret2",
+            };
+
+            try
+            {
+                options.HistorySavePath = newHistoryFilePath;
+                options.HistorySaveStyle = newHistorySaveStyle;
+                SetHistory(expectedHistoryItems);
+
+                // Sensitive input history should be kept in the internal history queue.
+                var historyItems = PSConsoleReadLine.GetHistoryItems();
+                Assert.Equal(expectedHistoryItems.Length, historyItems.Length);
+                for (int i = 0; i < expectedHistoryItems.Length; i++)
+                {
+                    Assert.Equal(expectedHistoryItems[i], historyItems[i].CommandLine);
+                }
+
+                // Sensitive input history should NOT be saved to the history file.
+                string[] text = File.ReadAllLines(newHistoryFilePath);
+                Assert.Equal(expectedSavedItems.Length, text.Length);
+                for (int i = 0; i < text.Length; i++)
+                {
+                    Assert.Equal(expectedSavedItems[i], text[i]);
+                }
+            }
+            finally
+            {
+                options.HistorySavePath = oldHistoryFilePath;
+                options.HistorySaveStyle = oldHistorySaveStyle;
+                File.Delete(newHistoryFilePath);
+            }
+        }
+
         [SkippableFact]
         public void SensitiveHistoryOptionalBehavior()
         {