Main.cs

﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Management;
using System.Reflection;
using System.Runtime.InteropServices;
using System.ServiceProcess;
using System.Text;
using System.Threading;
using Microsoft.Win32;
using WMI;
using ServiceType = WMI.ServiceType;

namespace winsw
{
    public class WrapperService : ServiceBase, EventLogger
    {
        private readonly ServiceDescriptor _descriptor;
        private readonly Process _process = new Process();
        private Dictionary<string, string> _envs;

        /// <summary>
        ///     Indicates to the watch dog thread that we are going to terminate the process,
        ///     so don't try to kill us when the child exits.
        /// </summary>
        private bool _orderlyShutdown;

        private bool _systemShuttingdown;
        private SERVICE_STATUS _wrapperServiceStatus;

        public WrapperService(ServiceDescriptor descriptor)
        {
            _descriptor = descriptor;
            ServiceName = _descriptor.Id;
            CanShutdown = true;
            CanStop = true;
            CanPauseAndContinue = false;
            AutoLog = true;
            _systemShuttingdown = false;
        }

        public WrapperService() : this(new ServiceDescriptor())
        {
        }

        /// <summary>
        ///     Version of Windows service wrapper
        /// </summary>
        /// <remarks>
        ///     The version will be taken from <see cref="AssemblyInfo" />
        /// </remarks>
        public static Version Version
        {
            get { return Assembly.GetExecutingAssembly().GetName().Version; }
        }

        public void LogEvent(string message)
        {
            if (_systemShuttingdown)
            {
                /* NOP - cannot call EventLog because of shutdown. */
            }
            else
            {
                try
                {
                    EventLog.WriteEntry(message);
                }
                catch (Exception e)
                {
                    WriteEvent("Failed to log event in Windows Event Log: " + message + "; Reason: ", e);
                }
            }
        }

        public void LogEvent(string message, EventLogEntryType type)
        {
            if (_systemShuttingdown)
            {
                /* NOP - cannot call EventLog because of shutdown. */
            }
            else
            {
                try
                {
                    EventLog.WriteEntry(message, type);
                }
                catch (Exception e)
                {
                    WriteEvent("Failed to log event in Windows Event Log. Reason: ", e);
                }
            }
        }

        /// <summary>
        ///     Process the file copy instructions, so that we can replace files that are always in use while
        ///     the service runs.
        /// </summary>
        private void HandleFileCopies()
        {
            var file = _descriptor.BasePath + ".copies";
            if (!File.Exists(file))
                return; // nothing to handle

            try
            {
                using (var tr = new StreamReader(file, Encoding.UTF8))
                {
                    string line;
                    while ((line = tr.ReadLine()) != null)
                    {
                        LogEvent("Handling copy: " + line);
                        var tokens = line.Split('>');
                        if (tokens.Length > 2)
                        {
                            LogEvent("Too many delimiters in " + line);
                            continue;
                        }

                        CopyFile(tokens[0], tokens[1]);
                    }
                }
            }
            finally
            {
                File.Delete(file);
            }
        }

        /// <summary>
        ///     File replacement.
        /// </summary>
        private void CopyFile(string sourceFileName, string destFileName)
        {
            try
            {
                File.Delete(destFileName);
                File.Move(sourceFileName, destFileName);
            }
            catch (IOException e)
            {
                LogEvent("Failed to copy :" + sourceFileName + " to " + destFileName + " because " + e.Message);
            }
        }

        /// <summary>
        ///     Starts a thread that protects the execution with a try/catch block.
        ///     It appears that in .NET, unhandled exception in any thread causes the app to terminate
        ///     http://msdn.microsoft.com/en-us/library/ms228965.aspx
        /// </summary>
        private void StartThread(ThreadStart main)
        {
            new Thread(delegate()
            {
                try
                {
                    main();
                }
                catch (Exception e)
                {
                    WriteEvent("Thread failed unexpectedly", e);
                }
            }).Start();
        }

        /// <summary>
        ///     Handle the creation of the logfiles based on the optional logmode setting.
        /// </summary>
        private void HandleLogfiles()
        {
            var logDirectory = _descriptor.LogDirectory;

            if (!Directory.Exists(logDirectory))
            {
                Directory.CreateDirectory(logDirectory);
            }

            var logAppender = _descriptor.LogHandler;
            logAppender.EventLogger = this;
            logAppender.log(_process.StandardOutput.BaseStream, _process.StandardError.BaseStream);
        }

        private void WriteEvent(Exception exception)
        {
            WriteEvent(exception.Message + "\nStacktrace:" + exception.StackTrace);
        }

        private void WriteEvent(string message, Exception exception)
        {
            WriteEvent(message + "\nMessage:" + exception.Message + "\nStacktrace:" + exception.StackTrace);
        }

        private void WriteEvent(string message)
        {
            var logfilename = Path.Combine(_descriptor.LogDirectory, _descriptor.BaseName + ".wrapper.log");
            var log = new StreamWriter(logfilename, true);

            log.WriteLine(DateTime.Now.ToString("yyyy-MM-dd HH:mm:ss") + " - " + message);
            log.Flush();
            log.Close();
        }

        protected override void OnStart(string[] _)
        {
            _envs = _descriptor.EnvironmentVariables;
            foreach (var key in _envs.Keys)
            {
                LogEvent("envar " + key + '=' + _envs[key]);
            }

            HandleFileCopies();

            // handle downloads
            foreach (var d in _descriptor.Downloads)
            {
                LogEvent("Downloading: " + d.From + " to " + d.To);
                try
                {
                    d.Perform();
                }
                catch (Exception e)
                {
                    LogEvent("Failed to download " + d.From + " to " + d.To + "\n" + e.Message);
                    WriteEvent("Failed to download " + d.From + " to " + d.To, e);
                    // but just keep going
                }
            }

            var startarguments = _descriptor.Startarguments;

            if (startarguments == null)
            {
                startarguments = _descriptor.Arguments;
            }
            else
            {
                startarguments += " " + _descriptor.Arguments;
            }

            LogEvent("Starting " + _descriptor.Executable + ' ' + startarguments);
            WriteEvent("Starting " + _descriptor.Executable + ' ' + startarguments);

            StartProcess(_process, startarguments, _descriptor.Executable);

            // send stdout and stderr to its respective output file.
            HandleLogfiles();

            _process.StandardInput.Close(); // nothing for you to read!
        }

        protected override void OnShutdown()
        {
//            WriteEvent("OnShutdown");

            try
            {
                _systemShuttingdown = true;
                StopIt();
            }
            catch (Exception ex)
            {
                WriteEvent("Shutdown exception", ex);
            }
        }

        protected override void OnStop()
        {
//            WriteEvent("OnStop");

            try
            {
                StopIt();
            }
            catch (Exception ex)
            {
                WriteEvent("Stop exception", ex);
            }
        }

        /// <summary>
        ///     Called when we are told by Windows SCM to exit.
        /// </summary>
        private void StopIt()
        {
            var stoparguments = _descriptor.Stoparguments;
            LogEvent("Stopping " + _descriptor.Id);
            WriteEvent("Stopping " + _descriptor.Id);
            _orderlyShutdown = true;

            if (stoparguments == null)
            {
                try
                {
                    WriteEvent("ProcessKill " + _process.Id);
                    StopProcessAndChildren(_process.Id);
                }
                catch (InvalidOperationException)
                {
                    // already terminated
                }
            }
            else
            {
                SignalShutdownPending();

                stoparguments += " " + _descriptor.Arguments;

                var stopProcess = new Process();
                var executable = _descriptor.StopExecutable;

                if (executable == null)
                {
                    executable = _descriptor.Executable;
                }

                StartProcess(stopProcess, stoparguments, executable);

                WriteEvent("WaitForProcessToExit " + _process.Id + "+" + stopProcess.Id);
                WaitForProcessToExit(_process);
                WaitForProcessToExit(stopProcess);
                SignalShutdownComplete();
            }

            if (_systemShuttingdown && _descriptor.BeepOnShutdown)
            {
                Console.Beep();
            }

            WriteEvent("Finished " + _descriptor.Id);
        }

        private void StopProcessAndChildren(int pid)
        {
            var childPids = GetChildPids(pid);

            if (_descriptor.StopParentProcessFirst)
            {
                StopProcess(pid);
                foreach (var childPid in childPids)
                {
                    StopProcessAndChildren(childPid);
                }
            }
            else
            {
                foreach (var childPid in childPids)
                {
                    StopProcessAndChildren(childPid);
                }
                StopProcess(pid);
            }
        }

        private List<int> GetChildPids(int pid)
        {
            var searcher = new ManagementObjectSearcher("Select * From Win32_Process Where ParentProcessID=" + pid);
            var childPids = new List<int>();
            foreach (var mo in searcher.Get())
            {
                var childProcessId = mo["ProcessID"];
                WriteEvent("Found child process: " + childProcessId + " Name: " + mo["Name"]);
                childPids.Add(Convert.ToInt32(childProcessId));
            }
            return childPids;
        }

        private void StopProcess(int pid)
        {
            WriteEvent("Stopping process " + pid);
            Process proc;
            try
            {
                proc = Process.GetProcessById(pid);
            }
            catch (ArgumentException)
            {
                WriteEvent("Process " + pid + " is already stopped");
                return;
            }

            WriteEvent("Send SIGINT " + pid);
            var successful = SigIntHelper.SendSIGINTToProcess(proc, _descriptor.StopTimeout);
            if (successful)
            {
                WriteEvent("SIGINT to" + pid + " successful");
            }
            else
            {
                try
                {
                    WriteEvent("SIGINT to " + pid + " failed - Killing as fallback");
                    proc.Kill();
                }
                catch (ArgumentException)
                {
                    // Process already exited.
                }
            }
        }

        private void WaitForProcessToExit(Process processoWait)
        {
            SignalShutdownPending();

            try
            {
//                WriteEvent("WaitForProcessToExit [start]");

                while (!processoWait.WaitForExit(_descriptor.SleepTime.Milliseconds))
                {
                    SignalShutdownPending();
//                    WriteEvent("WaitForProcessToExit [repeat]");
                }
            }
            catch (InvalidOperationException)
            {
                // already terminated
            }

//            WriteEvent("WaitForProcessToExit [finished]");
        }

        private void SignalShutdownPending()
        {
            var handle = ServiceHandle;
            _wrapperServiceStatus.checkPoint++;
            _wrapperServiceStatus.waitHint = _descriptor.WaitHint.Milliseconds;
//            WriteEvent("SignalShutdownPending " + wrapperServiceStatus.checkPoint + ":" + wrapperServiceStatus.waitHint);
            _wrapperServiceStatus.currentState = (int) State.SERVICE_STOP_PENDING;
            Advapi32.SetServiceStatus(handle, ref _wrapperServiceStatus);
        }

        private void SignalShutdownComplete()
        {
            var handle = ServiceHandle;
            _wrapperServiceStatus.checkPoint++;
//            WriteEvent("SignalShutdownComplete " + wrapperServiceStatus.checkPoint + ":" + wrapperServiceStatus.waitHint);
            _wrapperServiceStatus.currentState = (int) State.SERVICE_STOPPED;
            Advapi32.SetServiceStatus(handle, ref _wrapperServiceStatus);
        }

        private void StartProcess(Process processToStart, string arguments, string executable)
        {
            var ps = processToStart.StartInfo;
            ps.FileName = executable;
            ps.Arguments = arguments;
            ps.WorkingDirectory = _descriptor.WorkingDirectory;
            ps.CreateNoWindow = false;
            ps.UseShellExecute = false;
            ps.RedirectStandardInput = true;
                // this creates a pipe for stdin to the new process, instead of having it inherit our stdin.
            ps.RedirectStandardOutput = true;
            ps.RedirectStandardError = true;

            foreach (var key in _envs.Keys)
                Environment.SetEnvironmentVariable(key, _envs[key]);
            // ps.EnvironmentVariables[key] = envs[key]; // bugged (lower cases all variable names due to StringDictionary being used, see http://connect.microsoft.com/VisualStudio/feedback/ViewFeedback.aspx?FeedbackID=326163)

            processToStart.Start();
            WriteEvent("Started " + processToStart.Id);

            var priority = _descriptor.Priority;
            if (priority != ProcessPriorityClass.Normal)
                processToStart.PriorityClass = priority;

            // monitor the completion of the process
            StartThread(delegate
            {
                var msg = processToStart.Id + " - " + processToStart.StartInfo.FileName + " " +
                          processToStart.StartInfo.Arguments;
                processToStart.WaitForExit();

                try
                {
                    if (_orderlyShutdown)
                    {
                        LogEvent("Child process [" + msg + "] terminated with " + processToStart.ExitCode,
                            EventLogEntryType.Information);
                    }
                    else
                    {
                        LogEvent("Child process [" + msg + "] finished with " + processToStart.ExitCode,
                            EventLogEntryType.Warning);
                        // if we finished orderly, report that to SCM.
                        // by not reporting unclean shutdown, we let Windows SCM to decide if it wants to
                        // restart the service automatically
                        if (processToStart.ExitCode == 0)
                            SignalShutdownComplete();
                        Environment.Exit(processToStart.ExitCode);
                    }
                }
                catch (InvalidOperationException ioe)
                {
                    LogEvent("WaitForExit " + ioe.Message);
                }

                try
                {
                    processToStart.Dispose();
                }
                catch (InvalidOperationException ioe)
                {
                    LogEvent("Dispose " + ioe.Message);
                }
            });
        }

        public static int Main(string[] args)
        {
            try
            {
                Run(args);
                return 0;
            }
            catch (WmiException e)
            {
                Console.Error.WriteLine(e);
                return (int) e.ErrorCode;
            }
            catch (Exception e)
            {
                Console.Error.WriteLine(e);
                return -1;
            }
        }

        private static void ThrowNoSuchService()
        {
            throw new WmiException(ReturnValue.NoSuchService);
        }

        // ReSharper disable once InconsistentNaming
        public static void Run(string[] _args, ServiceDescriptor descriptor = null)
        {
            if (_args.Length > 0)
            {
                var d = descriptor ?? new ServiceDescriptor();
                var svc = new WmiRoot().GetCollection<Win32Services>();
                var s = svc.Select(d.Id);

                var args = new List<string>(Array.AsReadOnly(_args));
                if (args[0] == "/redirect")
                {
                    // Redirect output
                    // One might ask why we support this when the caller 
                    // can redirect the output easily. The answer is for supporting UAC.
                    // On UAC-enabled Windows such as Vista, SCM operation requires
                    // elevated privileges, thus winsw.exe needs to be launched
                    // accordingly. This in turn limits what the caller can do,
                    // and among other things it makes it difficult for the caller
                    // to read stdout/stderr. Thus redirection becomes handy.
                    var f = new FileStream(args[1], FileMode.Create);
                    var w = new StreamWriter(f) {AutoFlush = true};
                    Console.SetOut(w);
                    Console.SetError(w);

                    var handle = f.Handle;
                    Kernel32.SetStdHandle(-11, handle); // set stdout
                    Kernel32.SetStdHandle(-12, handle); // set stder

                    args = args.GetRange(2, args.Count - 2);
                }

                args[0] = args[0].ToLower();
                if (args[0] == "install")
                {
                    // Check if the service exists
                    if (s != null)
                    {
                        Console.WriteLine("Service with id '" + d.Id + "' already exists");
                        Console.WriteLine(
                            "To install the service, delete the existing one or change service Id in the configuration file");
                        throw new Exception("Installation failure: Service with id '" + d.Id + "' already exists");
                    }

                    string username = null, password = null;
                    var setallowlogonasaserviceright = false;
                    if (args.Count > 1 && args[1] == "/p")
                    {
                        // we expected username/password on stdin
                        Console.Write("Username: ");
                        username = Console.ReadLine();
                        Console.Write("Password: ");
                        password = ReadPassword();
                        Console.WriteLine();
                        Console.Write("Set Account rights to allow log on as a service (y/n)?: ");
                        var keypressed = Console.ReadKey();
                        Console.WriteLine();
                        if (keypressed.Key == ConsoleKey.Y)
                        {
                            setallowlogonasaserviceright = true;
                        }
                    }
                    else
                    {
                        if (d.HasServiceAccount())
                        {
                            username = d.ServiceAccountUser;
                            password = d.ServiceAccountPassword;
                            setallowlogonasaserviceright = d.AllowServiceAcountLogonRight;
                        }
                    }

                    if (setallowlogonasaserviceright)
                    {
                        LogonAsAService.AddLogonAsAServiceRight(username);
                    }

                    svc.Create(
                        d.Id,
                        d.Caption,
                        "\"" + d.ExecutablePath + "\"",
                        ServiceType.OwnProcess,
                        ErrorControl.UserNotified,
                        d.StartMode,
                        d.Interactive,
                        username,
                        password,
                        d.ServiceDependencies);

                    // update the description
                    /* Somehow this doesn't work, even though it doesn't report an error
                    Win32Service s = svc.Select(d.Id);
                    s.Description = d.Description;
                    s.Commit();
                     */

                    // so using a classic method to set the description. Ugly.
                    Registry.LocalMachine.OpenSubKey("System").OpenSubKey("CurrentControlSet").OpenSubKey("Services")
                        .OpenSubKey(d.Id, true).SetValue("Description", d.Description);

                    var actions = d.FailureActions;
                    if (actions.Count > 0)
                    {
// set the failure actions
                        using (var scm = new ServiceManager())
                        {
                            using (var sc = scm.Open(d.Id))
                            {
                                sc.ChangeConfig(d.ResetFailureAfter, actions);
                            }
                        }
                    }
                    return;
                }
                if (args[0] == "uninstall")
                {
                    if (s == null)
                    {
                        Console.WriteLine("Warning! The service with id '" + d.Id +
                                          "' does not exist. Nothing to uninstall");
                        return; // there's no such service, so consider it already uninstalled
                    }
                    try
                    {
                        s.Delete();
                    }
                    catch (WmiException e)
                    {
                        if (e.ErrorCode == ReturnValue.ServiceMarkedForDeletion)
                            return; // it's already uninstalled, so consider it a success
                        throw e;
                    }
                    return;
                }
                if (args[0] == "start")
                {
                    if (s == null) ThrowNoSuchService();
                    s.StartService();
                    return;
                }
                if (args[0] == "stop")
                {
                    if (s == null) ThrowNoSuchService();
                    s.StopService();
                    return;
                }
                if (args[0] == "restart")
                {
                    if (s == null)
                        ThrowNoSuchService();

                    if (s.Started)
                        s.StopService();

                    while (s.Started)
                    {
                        Thread.Sleep(1000);
                        s = svc.Select(d.Id);
                    }

                    s.StartService();
                    return;
                }
                if (args[0] == "restart!")
                {
                    // run restart from another process group. see README.md for why this is useful.

                    var si = new STARTUPINFO();
                    var pi = new PROCESS_INFORMATION();

                    var result = Kernel32.CreateProcess(null, d.ExecutablePath + " restart", IntPtr.Zero, IntPtr.Zero,
                        false, 0x200 /*CREATE_NEW_PROCESS_GROUP*/, IntPtr.Zero, null, ref si, out pi);
                    if (!result)
                    {
                        throw new Exception("Failed to invoke restart: " + Marshal.GetLastWin32Error());
                    }
                    return;
                }
                if (args[0] == "status")
                {
                    if (s == null)
                        Console.WriteLine("NonExistent");
                    else if (s.Started)
                        Console.WriteLine("Started");
                    else
                        Console.WriteLine("Stopped");
                    return;
                }
                if (args[0] == "test")
                {
                    var wsvc = new WrapperService();
                    wsvc.OnStart(args.ToArray());
                    Thread.Sleep(1000);
                    wsvc.OnStop();
                    return;
                }
                if (args[0] == "help" || args[0] == "--help" || args[0] == "-h"
                    || args[0] == "-?" || args[0] == "/?")
                {
                    printHelp();
                    return;
                }
                if (args[0] == "version")
                {
                    printVersion();
                    return;
                }

                Console.WriteLine("Unknown command: " + args[0]);
                printAvailableCommandsInfo();
                throw new Exception("Unknown command: " + args[0]);
            }
            Run(new WrapperService());
        }

        private static string ReadPassword()
        {
            var buf = new StringBuilder();
            ConsoleKeyInfo key;
            while (true)
            {
                key = Console.ReadKey(true);
                if (key.Key == ConsoleKey.Enter)
                {
                    return buf.ToString();
                }
                if (key.Key == ConsoleKey.Backspace)
                {
                    buf.Remove(buf.Length - 1, 1);
                    Console.Write("\b \b");
                }
                else
                {
                    Console.Write('*');
                    buf.Append(key.KeyChar);
                }
            }
        }

        private static void printHelp()
        {
            Console.WriteLine("A wrapper binary that can be used to host executables as Windows services");
            Console.WriteLine("");
            Console.WriteLine("Usage: winsw [/redirect file] <command> [<args>]");
            Console.WriteLine("       Missing arguments trigger the service mode");
            Console.WriteLine("");
            printAvailableCommandsInfo();
            Console.WriteLine("");
            Console.WriteLine("Extra options:");
            Console.WriteLine("- '/redirect' - redirect the wrapper's STDOUT and STDERR to the specified file");
            Console.WriteLine("");
            printVersion();
            Console.WriteLine("More info: https://github.com/kohsuke/winsw");
            Console.WriteLine("Bug tracker: https://github.com/kohsuke/winsw/issues");
        }

        //TODO: Rework to enum in winsw-2.0
        private static void printAvailableCommandsInfo()
        {
            Console.WriteLine("Available commands:");
            Console.WriteLine("- 'install'   - install the service to Windows Service Controller");
            Console.WriteLine("- 'uninstall' - uninstall the service");
            Console.WriteLine("- 'start'     - start the service (must be installed before)");
            Console.WriteLine("- 'stop'      - stop the service");
            Console.WriteLine("- 'restart'   - restart the service");
            Console.WriteLine("- 'restart!'  - self-restart (can be called from child processes)");
            Console.WriteLine("- 'status'    - check the current status of the service");
            Console.WriteLine("- 'test'      - check if the service can be started and then stopped");
            Console.WriteLine("- 'version'   - print the version info");
            Console.WriteLine("- 'help'      - print the help info (aliases: -h,--help,-?,/?)");
        }

        private static void printVersion()
        {
            Console.WriteLine("WinSW " + Version);
        }
    }
}