using System;
using System.ComponentModel;
using System.Diagnostics;
using System.Runtime.InteropServices;

namespace Rsdn.WatchDog
{
    /// <summary>
    /// Summary description for ShutdownHelper.
    /// </summary>
    internal class ShutdownHelper
    {
        private ShutdownHelper()
        {
        }

        [Flags]
            public enum TokenAccessType : uint
        {
            TOKEN_ASSIGN_PRIMARY    = 0x0001,
            TOKEN_DUPLICATE         = 0x0002,
            TOKEN_IMPERSONATE       = 0x0004,
            TOKEN_QUERY             = 0x0008,
            TOKEN_QUERY_SOURCE      = 0x0010,
            TOKEN_ADJUST_PRIVILEGES = 0x0020,
            TOKEN_ADJUST_GROUPS     = 0x0040,
            TOKEN_ADJUST_DEFAULT    = 0x0080,
            TOKEN_ADJUST_SESSIONID  = 0x0100,

            TOKEN_ALL_ACCESS        = 
                0xF0000 | 
                TOKEN_ASSIGN_PRIMARY |
                TOKEN_DUPLICATE |
                TOKEN_IMPERSONATE |
                TOKEN_QUERY |
                TOKEN_QUERY_SOURCE |
                TOKEN_ADJUST_PRIVILEGES |
                TOKEN_ADJUST_GROUPS |
                TOKEN_ADJUST_DEFAULT |
                TOKEN_ADJUST_SESSIONID,

            TOKEN_READ    = 
                0xF0000      |
                TOKEN_QUERY,

            TOKEN_WRITE =
                0xF0000     |
                TOKEN_ADJUST_PRIVILEGES   |
                TOKEN_ADJUST_GROUPS       |
                TOKEN_ADJUST_DEFAULT,

            TOKEN_EXECUTE = 
                0xF0000
        }

        [DllImport("advapi32.dll", SetLastError=true)]
        private static extern bool OpenProcessToken(IntPtr hProcess,
            TokenAccessType dwDesiredAccess, out IntPtr hToken);

        [DllImport("advapi32.dll", SetLastError=true)]
        public static extern bool LookupPrivilegeName(
            string lpSystemName, 
            [In]ref LUID lpLuid, 
            [Out] char[] lpName, 
            ref int cbName);

        [DllImport("advapi32.dll", SetLastError=true)]
        public static extern bool LookupPrivilegeValue(string lpSystemName,
            string lpName, out LUID Luid);

        [StructLayout(LayoutKind.Sequential)]
            public struct LUID
        {
            public int LowPart;
            public int HighPart;
        }    

        private class Luid
        {
            private readonly LUID _luid;
            public Luid(LUID luid)
            {
                _luid = luid;
            }
            /*public UInt64 Value
            {
                get
                {
                    return (((UInt64)(_luid.HighPart)) << 32) | ((UInt64)_luid.LowPart);
                }
            }*/
            internal LUID GetNativeLUID()
            {
                return _luid;
            }
        }

        [StructLayout(LayoutKind.Sequential)]
            public struct LUID_AND_ATTRIBUTES 
        {
            public LUID Luid;
            public int Attributes;
        }

        [StructLayout(LayoutKind.Sequential)]
            public struct TOKEN_PRIVILEGES 
        {
            public const int Size = 1;
            public int PrivilegeCount;
            public LUID_AND_ATTRIBUTES Privileges;
        }

        [DllImport("advapi32.dll", SetLastError=true)]
        public static extern bool AdjustTokenPrivileges(
            IntPtr TokenHandle, 
            bool DisableAllPrivileges, 
            ref TOKEN_PRIVILEGES NewState, 
            int BufferLength, 
            IntPtr PreviousState, 
            out int ReturnLength);

        [Flags]
            public enum PrivilegeAttributes : uint
        {
            Disabled          = 0,
            EnabledByDefault  = 0x00000001,
            Enabled           = 0x00000002,
            UsedForAccess     = 0x80000000,
        }

        [DllImport("advapi32.dll", SetLastError = true)]
        private static extern bool InitiateSystemShutdown(
            [MarshalAs(UnmanagedType.LPTStr)]string machineName,
            [MarshalAs(UnmanagedType.LPTStr)]string message,
            int timeout, bool forceAppClosed, bool rebootAfterShutdown);

        public static void DoShutdown(string message)
        {
            Console.WriteLine("Restart");
            IntPtr token;
            if (!OpenProcessToken(Process.GetCurrentProcess().Handle,
                TokenAccessType.TOKEN_ADJUST_PRIVILEGES | TokenAccessType.TOKEN_QUERY,
                out token))
                ReportWin32Error();

            TOKEN_PRIVILEGES tkp = new TOKEN_PRIVILEGES();
            tkp.PrivilegeCount = TOKEN_PRIVILEGES.Size;
            tkp.Privileges.Attributes = 2;
            if (!LookupPrivilegeValue(null, "SeShutdownPrivilege", out tkp.Privileges.Luid))
                ReportWin32Error();
            
            int rl = 0;
            if (!AdjustTokenPrivileges(token, false, ref tkp, 0, IntPtr.Zero, out rl))
                ReportWin32Error();

            if (!InitiateSystemShutdown(null, message,
                0, true, true))
                ReportWin32Error();

            EventLog.WriteEntry(RsdnWD.Name, message, EventLogEntryType.Error);
        }

        private static void ReportWin32Error()
        {
            string msg = new Win32Exception().Message;
            Console.WriteLine(msg);
            EventLog.WriteEntry(RsdnWD.Name, msg,
                EventLogEntryType.Error);
        }
    }
}