Added reauth command to refresh existing agent settings with new cred…

…entials. (#4991)
microsoft · Sep 17, 2024 · 6b8de4d · 6b8de4d
1 parent 08552ff
commit 6b8de4d
Show file tree

Hide file tree

Showing 9 changed files with 381 additions and 61 deletions.
diff --git a/src/Agent.Listener/Agent.cs b/src/Agent.Listener/Agent.cs
@@ -128,6 +128,21 @@ public async Task<int> ExecuteCommand(CommandSettings command)
                     }
                 }
 
+                if (command.IsReAuthCommand())
+                {
+                    try
+                    {
+                        await configManager.ReAuthAsync(command);
+                        return Constants.Agent.ReturnCode.Success;
+                    }
+                    catch (Exception ex)
+                    {
+                        Trace.Error(ex);
+                        _term.WriteError(ex.Message);
+                        return Constants.Agent.ReturnCode.TerminatedError;
+                    }
+                }
+
                 _inConfigStage = false;
 
                 // warmup agent process (JIT/CLR)

diff --git a/src/Agent.Listener/CommandLine/ReAuthAgent.cs b/src/Agent.Listener/CommandLine/ReAuthAgent.cs
@@ -0,0 +1,10 @@
+using CommandLine;
+using Microsoft.VisualStudio.Services.Agent;
+
+namespace Agent.Listener.CommandLine
+{
+    [Verb(Constants.Agent.CommandLine.Commands.ReAuth)]
+    public class ReAuthAgent : ConfigureOrRemoveBase
+    {
+    }
+}
diff --git a/src/Agent.Listener/CommandSettings.cs b/src/Agent.Listener/CommandSettings.cs
@@ -30,6 +30,7 @@ public sealed class CommandSettings
             typeof(RunAgent),
             typeof(RemoveAgent),
             typeof(WarmupAgent),
+            typeof(ReAuthAgent),
         };
 
         private string[] verbCommands = new string[]
@@ -38,6 +39,7 @@ public sealed class CommandSettings
             Constants.Agent.CommandLine.Commands.Remove,
             Constants.Agent.CommandLine.Commands.Run,
             Constants.Agent.CommandLine.Commands.Warmup,
+            Constants.Agent.CommandLine.Commands.ReAuth,
         };
 
 
@@ -46,6 +48,7 @@ public sealed class CommandSettings
         private RemoveAgent Remove { get; set; }
         private RunAgent Run { get; set; }
         private WarmupAgent Warmup { get; set; }
+        private ReAuthAgent ReAuth { get; set; }
 
         public IEnumerable<Error> ParseErrors { get; set; }
 
@@ -86,6 +89,11 @@ public CommandSettings(IHostContext context, string[] args, IScopedEnvironment e
                 context.SecretMasker.AddValue(Remove.ClientSecret, WellKnownSecretAliases.RemoveClientSecret);
             }
 
+            if (ReAuth != null)
+            {
+                context.SecretMasker.AddValue(ReAuth.Token, WellKnownSecretAliases.RemoveToken);
+            }
+
             PrintArguments();
 
             // Store and remove any args passed via environment variables.
@@ -609,7 +617,8 @@ public bool IsVersion()
             if ((Configure?.Version == true) ||
                 (Remove?.Version == true) ||
                 (Run?.Version == true) ||
-                (Warmup?.Version == true))
+                (Warmup?.Version == true) ||
+                (ReAuth?.Version == true))
             {
                 return true;
             }
@@ -622,7 +631,8 @@ public bool IsHelp()
             if ((Configure?.Help == true) ||
                 (Remove?.Help == true) ||
                 (Run?.Help == true) ||
-                (Warmup?.Help == true))
+                (Warmup?.Help == true) ||
+                (ReAuth?.Help == true))
             {
                 return true;
             }
@@ -670,6 +680,8 @@ public bool IsWarmupCommand()
             return false;
         }
 
+        public bool IsReAuthCommand() => ReAuth != null;
+
 
         //
         // Private helpers.
@@ -839,6 +851,11 @@ private void ParseArguments(string[] args, bool ignoreErrors)
                         {
                             Warmup = x;
                         })
+                    .WithParsed<ReAuthAgent>(
+                        x =>
+                        {
+                            ReAuth = x;
+                        })
                     .WithNotParsed(
                         errors =>
                         {
@@ -859,6 +876,11 @@ private void PrintArguments()
                 _trace.Info(string.Concat(nameof(Remove), " ", ObjectAsJson(Remove)));
             }
 
+            if (ReAuth != null)
+            {
+                _trace.Info(string.Concat(nameof(ReAuth), " ", ObjectAsJson(ReAuth)));
+            }
+
             if (Warmup != null)
             {
                 _trace.Info(string.Concat(nameof(Warmup), " ", ObjectAsJson(Warmup)));
@@ -889,6 +911,11 @@ private ConfigureOrRemoveBase GetConfigureOrRemoveBase()
                 return Remove as ConfigureOrRemoveBase;
             }
 
+            if (ReAuth != null)
+            {
+                return ReAuth as ConfigureOrRemoveBase;
+            }
+
             return null;
         }
     }