SlunkCrypt/gui/Utilities/ProcessRunner.cs

/******************************************************************************/
/* SlunkCrypt, by LoRd_MuldeR <MuldeR2@GMX.de>                                */
/* This work has been released under the CC0 1.0 Universal license!           */
/******************************************************************************/

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Collections.Specialized;
using System.Diagnostics;
using System.IO;
using System.Text;
using System.Threading.Tasks;
using System.Windows.Threading;

namespace com.muldersoft.slunkcrypt.gui.utils
{
    class ProcessRunner : IDisposable
    {
        public delegate void OutputLineHandler(string line, bool stream);
        public delegate void ProgressChangedHandler(double progress);
        public delegate bool ProgressStringHandler(string line, out double progress);

        public event OutputLineHandler OutputAvailable;
        public event ProgressChangedHandler ProgressChanged;
        public const bool STDOUT = false, STDERR = true;

        private readonly Process m_process = new Process();
        private readonly Dispatcher m_dispatcher;
        private readonly ProcessPriorityClass? m_priorityClass;
        private readonly Task<int> m_hasExited;

        private volatile bool m_running = false, m_finished = false, m_aborted = false, m_disposed = false;

        // =============================================================================
        // Exception classes
        // =============================================================================

        public class ProcessStartException : IOException
        {
            public ProcessStartException(string message, Exception innerException) : base(message, innerException)
            {
            }
        }

        public class ProcessInterruptedException : IOException
        {
            public ProcessInterruptedException(string message) : base(message)
            {
            }
        }

        // =============================================================================
        // Event handler class
        // =============================================================================

        private class ProcessExitHandler
        {
            private readonly TaskCompletionSource<int> completionSource = new TaskCompletionSource<int>();

            public ProcessExitHandler(Process process)
            {
                if (ReferenceEquals(process, null))
                {
                    throw new ArgumentNullException("Process must not be null!");
                }
                process.Exited += OnProcessExit;
            }

            public Task<int> Task
            {
                get
                {
                    return completionSource.Task;
                }
            }

            protected void OnProcessExit(object sender, EventArgs e)
            {
                if (sender is Process)
                {
                    completionSource.TrySetResult(((Process)sender).ExitCode);
                }
            }
        }

        // =============================================================================
        // Constructor
        // =============================================================================

        public ProcessRunner(Dispatcher dispatcher, ProcessPriorityClass? priorityClass = null)
        {
            if (ReferenceEquals(m_dispatcher = dispatcher, null))
            {
                throw new ArgumentException("Dispatcher must not be null!");
            }
            if ((m_priorityClass = priorityClass).HasValue && (!Enum.IsDefined(typeof(ProcessPriorityClass), priorityClass.Value)))
            {
                throw new ArgumentException("The given ProcessPriorityClass is undefined!");
            }
            m_hasExited = InitializeProcess(m_process, true);
        }

        ~ProcessRunner()
        {
            Dispose(); /*just to be sure*/
        }

        // =============================================================================
        // Public methods
        // =============================================================================

        public async Task<int> ExecAsnyc(string executablePath, string[] arguments = null, string workingDirectory = null, IReadOnlyDictionary<string, string> environmentVariables = null)
        {
            m_dispatcher.VerifyAccess();
            if (m_disposed)
            {
                throw new ObjectDisposedException("ProcessRunner");
            }
            if (m_running || m_finished)
            {
                throw new InvalidOperationException("Process is still running or has already finished!");
            }
            m_running = true;
            try
            {
                return await DoExecAsnyc(executablePath, arguments, workingDirectory, environmentVariables);
            }
            finally
            {
                m_finished = true;
                m_running = false;
            }
        }

        public static async Task<Tuple<int, string[]>> ExecAsnyc(string executableFile, string[] arguments = null, string workingDirectory = null, IReadOnlyDictionary<string, string> environmentVariables = null, ProcessPriorityClass? priorityClass = null, TimeSpan? timeout = null)
        {
            
            using (Process process = new Process())
            {
                Task<int> hasExitedTask = InitializeProcess(process);
                return await DoExecAsnyc(process, hasExitedTask, executableFile, arguments, workingDirectory, environmentVariables, priorityClass, timeout);
            }
        }

        public void AbortProcess()
        {
            if ((!m_disposed) && m_running)
            {
                m_aborted = true;
                KillProcess();
            }
        }

        public virtual void Dispose()
        {
            if (!m_disposed)
            {
                m_disposed = true;
                GC.SuppressFinalize(this);
                KillProcess();
                try
                {
                    m_process.Dispose();
                }
                catch { }
            }
        }

        // =============================================================================
        // Protected methods
        // =============================================================================

        protected virtual double ParseProgressString(ref string currentLine, bool stream)
        {
            return double.NaN;
        }

        // =============================================================================
        // Internal methods
        // =============================================================================

        private static Task<int> InitializeProcess(Process process, bool redirStdErr = false)
        {
            process.StartInfo.UseShellExecute = false;
            process.StartInfo.CreateNoWindow = true;
            process.StartInfo.RedirectStandardOutput = true;
            process.StartInfo.StandardOutputEncoding = Encoding.UTF8;
            if (redirStdErr)
            {
                process.StartInfo.RedirectStandardError = true;
                process.StartInfo.StandardErrorEncoding = Encoding.UTF8;
            }
            process.EnableRaisingEvents = true;
            return new ProcessExitHandler(process).Task;
        }

        private async Task<int> DoExecAsnyc(string executablePath, string[] arguments, string workingDirectory, IReadOnlyDictionary<string, string> environmentVariables)
        {
            try
            {
                StartProcess(m_process, executablePath, arguments, workingDirectory, environmentVariables, m_priorityClass);
            }
            catch (Exception err)
            {
                throw new ProcessStartException("Failed to create the process!", err);
            }
            return await WaitForExit();
        }

        private static async Task<Tuple<int,string[]>> DoExecAsnyc(Process process, Task<int> hasExited, string executablePath, string[] arguments, string workingDirectory, IReadOnlyDictionary<string, string> environmentVariables, ProcessPriorityClass? priorityClass, TimeSpan? timeout)
        {
            try
            {
                StartProcess(process, executablePath, arguments, workingDirectory, environmentVariables, priorityClass);
            }
            catch (Exception err)
            {
                throw new ProcessStartException("Failed to create the process!", err);
            }
            string[] outputLines = await WaitForExit(process, hasExited, timeout.GetValueOrDefault(TimeSpan.MaxValue));
            return Tuple.Create(hasExited.Result, outputLines);
        }

        private static void StartProcess(Process process, string executablePath, string[] arguments, string workingDirectory, IReadOnlyDictionary<string, string> environmentVariables, ProcessPriorityClass? priorityClass)
        {
            process.StartInfo.FileName = executablePath;
            process.StartInfo.Arguments = CreateArgumentList(arguments);
            process.StartInfo.WorkingDirectory = string.IsNullOrEmpty(workingDirectory) ? GetWorkingDirectory(executablePath) : workingDirectory;
            SetupEnvironment(process.StartInfo.EnvironmentVariables, executablePath, environmentVariables);
            process.Start();
            SetProcessPriority(process, priorityClass);
        }

        private async Task<int> WaitForExit()
        {
            Task readStdOutTask = Task.Run(() => ReadProcessOutput(m_process.StandardOutput, (line) => HandleOutput(line, STDOUT)));
            Task readStdErrTask = Task.Run(() => ReadProcessOutput(m_process.StandardError,  (line) => HandleOutput(line, STDERR)));
            await Task.WhenAll(readStdOutTask, readStdErrTask, m_hasExited);
            if (m_aborted || m_disposed)
            {
                NotifyOutputAvailable("\u2192 Process has been aborted !!!", true);
                throw new ProcessInterruptedException("Process has been aborted!");
            }
            int processExitCode = m_hasExited.Result;
            NotifyOutputAvailable(string.Format("\u2192 Process terminated normally [Exit status: {0:D}]", processExitCode), false);
            return processExitCode;
        }

        private static async Task<string[]> WaitForExit(Process process, Task<int> hasExited, TimeSpan timeout)
        {
            ConcurrentStack<string> outputLines = new ConcurrentStack<string>();
            Task readerTask = Task.Run(() => ReadProcessOutput(process.StandardOutput, (line) => outputLines.Push(line)));
            if (await Task.WhenAny(hasExited, Task.Delay(timeout)) != hasExited)
            {
                KillProcess(process);
            }
            await Task.WhenAll(readerTask, hasExited);
            return outputLines.ToArray();
        }

        private static void ReadProcessOutput(StreamReader reader, Action<string> outputHandler)
        {
            using (reader)
            {
                char[] buffer = new char[1024];
                StringBuilder lineBuffer = new StringBuilder();
                while (!reader.EndOfStream)
                {
                    ReadNextChunk(reader, buffer, lineBuffer, outputHandler);
                }
                ProcessLine(lineBuffer, outputHandler);
            }
        }

        private static void ReadNextChunk(StreamReader reader, char[] buffer, StringBuilder lineBuffer, Action<string> outputHandler)
        {
            int count = reader.Read(buffer, 0, buffer.Length);
            if (count > 0)
            {
                for (int i = 0; i < count; ++i)
                {
                    char c = buffer[i];
                    if ((c != '\0') && (c != '\n') && (c != '\r') && (c != '\b'))
                    {
                        lineBuffer.Append(c);
                    }
                    else
                    {
                        ProcessLine(lineBuffer, outputHandler);
                        lineBuffer.Clear();
                    }
                }
            }
        }

        private static void ProcessLine(StringBuilder lineBuffer, Action<string> outputHandler)
        {
            if (lineBuffer.Length > 0)
            {
                String currentLine;
                if (!string.IsNullOrWhiteSpace(currentLine = lineBuffer.ToString()))
                {
                    outputHandler.Invoke(currentLine.Trim());
                }
            }
        }

        private void HandleOutput(string currentLine, bool stream)
        {
            CheckForProgressUpdate(ref currentLine, stream);
            if (!String.IsNullOrEmpty(currentLine))
            {
                NotifyOutputAvailable(currentLine, stream);
            }
        }

        private void CheckForProgressUpdate(ref string currentLine, bool stream)
        {
            double progress = ParseProgressString(ref currentLine, stream);
            if (!double.IsNaN(progress))
            {
                NotifyProgressChanged(progress);
            }
        }

        private void NotifyOutputAvailable(string line, bool stream)
        {
            if (!ReferenceEquals(OutputAvailable, null))
            {
                m_dispatcher.InvokeAsync(() => OutputAvailable(line, stream));
            }
        }

        private void NotifyProgressChanged(double progress)
        {
            if (!ReferenceEquals(ProgressChanged, null))
            {
                m_dispatcher.InvokeAsync(() => ProgressChanged(Math.Max(0.0, Math.Min(1.0, progress))));
            }
        }

        private static string CreateArgumentList(string[] arguments)
        {
            StringBuilder argumentList = new StringBuilder();
            if (!ReferenceEquals(arguments, null))
            {
                for (int i = 0; i < arguments.Length; ++i)
                {
                    if (i > 0)
                    {
                        argumentList.Append("\x20");
                    }
                    argumentList.Append(EscapeArgument(arguments[i]));
                }
            }
            return argumentList.ToString();
        }

        private static string EscapeArgument(string str)
        {
            if (RequiresQuotation(str))
            {
                StringBuilder sb = new StringBuilder().Append('"');
                foreach (char c in str)
                {
                    if (c == '"')
                    {
                        sb.Append('\\').Append('"');
                    }
                    else
                    {
                        sb.Append(c);
                    }
                }
                return sb.Append('"').ToString();
            }
            return str;
        }

        private static bool RequiresQuotation(string str)
        {
            foreach (char c in str)
            {
                if (char.IsWhiteSpace(c) || (c == '"'))
                {
                    return true;
                }
            }
            return false;
        }

        private static string GetWorkingDirectory(string executablePath)
        {
            try
            {
                string directory = Path.GetDirectoryName(executablePath);
                if (!string.IsNullOrWhiteSpace(directory))
                {
                    return directory;
                }
            }
            catch { }
            return AppDomain.CurrentDomain.BaseDirectory;
        }

        private static void SetupEnvironment(StringDictionary dictionary, string executablePath, IReadOnlyDictionary<string, string> environmentVariables)
        {
            dictionary["PATH"] = PathUtils.CreatePathSpec(Path.GetDirectoryName(executablePath));
            if (!ReferenceEquals(environmentVariables, null))
            {
                foreach (KeyValuePair<string, string> entry in environmentVariables)
                {
                    dictionary[entry.Key] = entry.Value;
                }
            }
        }

        private void KillProcess()
        {
            try
            {
                m_process.Kill();
            }
            catch { }
        }

        private static void KillProcess(Process process)
        {
            try
            {
                process.Kill();
            }
            catch { }
        }

        private static void SetProcessPriority(Process process, ProcessPriorityClass? priorityClass)
        {
            try
            {
                if (priorityClass.HasValue)
                {
                    process.PriorityClass = priorityClass.Value;
                }
            }
            catch { }
        }
    }
}