import ctypes
import keystone
import enum

class Page(enum.IntEnum):
    EXECUTE           = 0x10
    EXECUTE_READ      = 0x20
    EXECUTE_READWRITE = 0x40
    EXECUTE_WRITECOPY = 0x80
    NOACCESS          = 0x01
    READONLY          = 0x02
    READWRITE         = 0x04
    WRITECOPY         = 0x08

class Memory(enum.IntFlag):
    COMMIT     = 0x00001000
    RESERVE    = 0x00002000
    RESET      = 0x00080000
    RESET_UNDO = 0x1000000
    DECOMMIT   = 0x00004000
    RELEASE    = 0x00008000

def win32_bool_ptr_errcheck(result, func, args):
    if not result:
        raise ctypes.WinError()
    return result

VirtualProtect = ctypes.windll.kernel32.VirtualProtect
VirtualProtect.restype = bool
VirtualProtect.argtypes = [ctypes.c_void_p, ctypes.c_size_t,
                           ctypes.c_int, ctypes.POINTER(ctypes.c_int)]
VirtualProtect.errcheck = win32_bool_ptr_errcheck

VirtualAlloc = ctypes.windll.kernel32.VirtualAlloc
VirtualAlloc.restype = ctypes.c_void_p
VirtualAlloc.argtypes = [ctypes.c_void_p, ctypes.c_size_t,
                         ctypes.c_int, ctypes.c_int]
VirtualAlloc.errcheck = win32_bool_ptr_errcheck

VirtualFree = ctypes.windll.kernel32.VirtualFree
VirtualFree.restype = bool
VirtualFree.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
VirtualFree.errcheck = win32_bool_ptr_errcheck

SQUARE_PROC = ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int)

CODE = b"""
mov rax, rcx
imul rax, rax
ret 0
"""

if __name__ == "__main__":
    ks = keystone.Ks(keystone.KS_ARCH_X86, keystone.KS_MODE_64)
    encoding, _ = ks.asm(CODE)

    memory = VirtualAlloc(None, len(encoding), Memory.COMMIT, Page.READWRITE)

    old_protect = ctypes.c_int(0)
    ctypes.memmove(memory, bytes(encoding), len(encoding))

    VirtualProtect(memory, len(encoding), Page.EXECUTE, ctypes.byref(old_protect))

    square = SQUARE_PROC(memory)
    print(square(2))

    VirtualFree(memory, 0, Memory.RELEASE)