src/AWS.Deploy.CLI/Utilities/CommandLineWrapper.cs

// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using Amazon.Runtime;
using AWS.Deploy.Common;
using AWS.Deploy.Orchestration;
using AWS.Deploy.Orchestration.Utilities;

namespace AWS.Deploy.CLI.Utilities
{
    public class CommandLineWrapper : ICommandLineWrapper
    {
        private readonly IOrchestratorInteractiveService _interactiveService;
        private readonly bool _useSeparateWindow;
        private Action<ProcessStartInfo>? _processStartInfoAction;

        public CommandLineWrapper(
            IOrchestratorInteractiveService interactiveService)
        {
            _interactiveService = interactiveService;
        }

        public CommandLineWrapper(
            IOrchestratorInteractiveService interactiveService,
            bool useSeparateWindow)
            : this(interactiveService)
        {
            _useSeparateWindow = useSeparateWindow;
        }

        /// <inheritdoc />
        public async Task Run(
            string command,
            string workingDirectory = "",
            bool streamOutputToInteractiveService = true,
            Action<TryRunResult>? onComplete = null,
            bool redirectIO = true,
            string? stdin = null,
            IDictionary<string, string>? environmentVariables = null,
            bool needAwsCredentials = false,
            CancellationToken cancellationToken = default)
        {
            StringBuilder strOutput = new StringBuilder();
            StringBuilder strError = new StringBuilder();

            var processStartInfo = new ProcessStartInfo
            {
                FileName = GetSystemShell(),

                Arguments =
                    RuntimeInformation.IsOSPlatform(OSPlatform.Windows)
                        ? $"/c {command}"
                        : $"-c \"{command}\"",

                RedirectStandardInput = redirectIO || stdin != null,
                RedirectStandardOutput = redirectIO,
                RedirectStandardError = redirectIO,
                UseShellExecute = false,
                CreateNoWindow = redirectIO,
                WorkingDirectory = workingDirectory
            };

            // If the command output is not being redirected check to see if
            // the output should go to a separate console window. This is important when run from
            // an IDE which won't have a console window by default.
            if (!streamOutputToInteractiveService && !redirectIO)
            {
                processStartInfo.UseShellExecute = _useSeparateWindow;
            }

            UpdateEnvironmentVariables(processStartInfo, environmentVariables);

            if (needAwsCredentials)
                _processStartInfoAction?.Invoke(processStartInfo);

            var process = Process.Start(processStartInfo);
            if (null == process)
                throw new Exception("Process.Start failed to return a non-null process");

            if (redirectIO)
            {
                process.OutputDataReceived += (sender, e) =>
                {
                    if(streamOutputToInteractiveService)
                    {
                        _interactiveService.LogInfoMessage(e.Data);
                    }

                    strOutput.AppendLine(e.Data);
                };

                process.ErrorDataReceived += (sender, e) =>
                {
                    if(streamOutputToInteractiveService)
                    {
                        _interactiveService.LogInfoMessage(e.Data);
                    }

                    strError.AppendLine(e.Data);
                };
                process.BeginOutputReadLine();
                process.BeginErrorReadLine();
            }

            if(stdin != null)
            {
                process.StandardInput.Write(stdin);
                process.StandardInput.Close();
            }

            await process.WaitForExitAsync(cancellationToken);

            if (onComplete != null)
            {
                var result = new TryRunResult
                {
                    ExitCode = process.ExitCode
                };

                if (redirectIO)
                {
                    result.StandardError = strError.ToString();
                    result.StandardOut = strOutput.ToString();
                }

                onComplete(result);
            }
        }

        private static void UpdateEnvironmentVariables(ProcessStartInfo processStartInfo, IDictionary<string, string>? environmentVariables)
        {
            if (environmentVariables == null)
            {
                return;
            }

            foreach (var (key, value) in environmentVariables)
            {
                if (key == EnvironmentVariableKeys.AWS_EXECUTION_ENV)
                {
                    var awsExecutionEnvValue = BuildAWSExecutionEnvValue(processStartInfo, value);
                    processStartInfo.EnvironmentVariables[key] = awsExecutionEnvValue;
                }
                else
                {
                    processStartInfo.EnvironmentVariables[key] = value;
                }
            }
        }

        private static string BuildAWSExecutionEnvValue(ProcessStartInfo processStartInfo, string awsExecutionEnv)
        {
            var awsExecutionEnvBuilder = new StringBuilder();
            if (processStartInfo.EnvironmentVariables.ContainsKey(EnvironmentVariableKeys.AWS_EXECUTION_ENV)
                && !string.IsNullOrEmpty(processStartInfo.EnvironmentVariables[EnvironmentVariableKeys.AWS_EXECUTION_ENV]))
            {
                awsExecutionEnvBuilder.Append(processStartInfo.EnvironmentVariables[EnvironmentVariableKeys.AWS_EXECUTION_ENV]);
            }

            if (!string.IsNullOrEmpty(awsExecutionEnv))
            {
                if (awsExecutionEnvBuilder.Length != 0)
                {
                    awsExecutionEnvBuilder.Append("_");
                }

                awsExecutionEnvBuilder.Append(awsExecutionEnv);
            }

            return awsExecutionEnvBuilder.ToString();
        }

        public void ConfigureProcess(Action<ProcessStartInfo>? processStartInfoAction)
        {
            _processStartInfoAction = processStartInfoAction;
        }

        private string GetSystemShell()
        {
            if (TryGetEnvironmentVariable("COMSPEC", out var comspec))
                return comspec!;

            if (TryGetEnvironmentVariable("SHELL", out var shell))
                return shell!;

            // fall back to defaults
            return RuntimeInformation.IsOSPlatform(OSPlatform.Windows)
                ? "cmd.exe"
                : "/bin/sh";
        }

        private bool TryGetEnvironmentVariable(string variable, out string? value)
        {
            value = Environment.GetEnvironmentVariable(variable);

            return !string.IsNullOrEmpty(value);
        }
    }
}