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.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 TaskCompletionSource<int> m_hasExited = new TaskCompletionSource<int>();
        private readonly Dispatcher m_dispatcher;
        private readonly ProcessPriorityClass? m_priorityClass;

        private volatile bool m_running = 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)
            {
            }
        }

        // =============================================================================
        // 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_process.StartInfo.UseShellExecute = false;
            m_process.StartInfo.CreateNoWindow  = true;
            m_process.StartInfo.RedirectStandardOutput = true;
            m_process.StartInfo.RedirectStandardError  = true;
            m_process.StartInfo.StandardOutputEncoding = Encoding.UTF8;
            m_process.StartInfo.StandardErrorEncoding  = Encoding.UTF8;
            m_process.EnableRaisingEvents = true;
            m_process.Exited += ProcessExitedEventHandler;
        }

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

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

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

        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(StringBuilder currentLine)
        {
            return double.NaN;
        }

        // =============================================================================
        // Event handlers
        // =============================================================================

        private void ProcessExitedEventHandler(object sender, EventArgs e)
        {
            if (m_process.HasExited)
            {
                m_hasExited.TrySetResult(m_process.ExitCode);
            }
        }

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

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

        private void StartProcess(string executablePath, string[] arguments, IReadOnlyDictionary<string, string> environmentVariables)
        {
            m_process.StartInfo.FileName = executablePath;
            m_process.StartInfo.Arguments = CreateArgumentList(arguments);
            SetupEnvironment(m_process.StartInfo.EnvironmentVariables, environmentVariables);
            m_process.StartInfo.WorkingDirectory = AppDomain.CurrentDomain.BaseDirectory;
            m_process.Start();
            SetProcessPriority(m_priorityClass);
        }

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

        private void ReadProcessOutput(StreamReader reader, bool stream)
        {
            using(reader)
            {
                char[] buffer = new char[1024];
                StringBuilder currentLine = new StringBuilder();
                while (!reader.EndOfStream)
                {
                    ReadNextChunk(reader, buffer, currentLine, stream);
                }
                NotifyOutputAvailable(currentLine, stream);
            }
        }

        private void ReadNextChunk(StreamReader reader, char[] buffer, StringBuilder currentLine, bool stderr)
        {
            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'))
                    {
                        currentLine.Append(c);
                    }
                    else
                    {
                        CheckForProgressUpdate(currentLine);
                        NotifyOutputAvailable(currentLine, stderr);
                        currentLine.Clear();
                    }
                }
                CheckForProgressUpdate(currentLine);
            }
        }

        private void CheckForProgressUpdate(StringBuilder currentLine)
        {
            if (currentLine.Length > 0)
            {
                double progress = ParseProgressString(currentLine);
                if (!(double.IsNaN(progress) || double.IsInfinity(progress)))
                {
                    NotifyProgressChanged(progress);
                }
            }
        }

        private void NotifyOutputAvailable(StringBuilder currentLine, bool stream)
        {
            if (currentLine.Length > 0)
            {
                String line = currentLine.ToString().Trim();
                if (!string.IsNullOrEmpty(line))
                {
                    NotifyOutputAvailable(line, stream);
                }
            }
        }

        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 void SetupEnvironment(StringDictionary dictionary, IReadOnlyDictionary<string, string> environmentVariables)
        {
            if (!ReferenceEquals(environmentVariables, null))
            {
                foreach (KeyValuePair<string, string> entry in environmentVariables)
                {
                    dictionary.Add(entry.Key, entry.Value);
                }
            }
        }

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

        private void SetProcessPriority(ProcessPriorityClass? priorityClass)
        {
            try
            {
                if (priorityClass.HasValue)
                {
                    m_process.PriorityClass = priorityClass.Value;
                }
            }
            catch { }
        }
    }
}