mirror of
https://github.com/hakril/PythonForWindows
synced 2026-06-08 14:31:45 +00:00
Clement Rouault
117 lines
3.3 KiB
Python
117 lines
3.3 KiB
Python
import ctypes
|
|
import mmap
|
|
import platform
|
|
import sys
|
|
|
|
import windows
|
|
import windows.winproxy
|
|
|
|
from . import simple_x86 as x86
|
|
from . import simple_x64 as x64
|
|
|
|
|
|
class PyObj(ctypes.Structure):
|
|
_fields_ = [("ob_refcnt", ctypes.c_size_t),
|
|
("ob_type", ctypes.c_void_p)] # must be cast
|
|
|
|
|
|
class PyMmap(PyObj):
|
|
_fields_ = [("ob_addr", ctypes.c_size_t), ("ob_size", ctypes.c_size_t)]
|
|
|
|
|
|
# Specific mmap class for code injection
|
|
class MyMap(mmap.mmap):
|
|
""" A mmap that is never unmapped and that contains the page address """
|
|
def __init__(self, *args, **kwarg):
|
|
# Get the page address by 'introspection' of the C struct
|
|
m = PyMmap.from_address(id(self))
|
|
self.addr = m.ob_addr
|
|
# Prevent garbage collection (so unmaping) of the page
|
|
m.ob_refcnt += 1
|
|
|
|
@classmethod
|
|
def get_map(cls, size):
|
|
""" Dispatch to the good mmap implem depending on the current system """
|
|
systems = {'windows': Win32MyMap,
|
|
'linux': UnixMyMap}
|
|
x = platform.system().lower()
|
|
if x not in systems:
|
|
raise ValueError("Unknow system {0}".format(x))
|
|
return systems[x].get_map(size)
|
|
|
|
|
|
class Win32MyMap(MyMap):
|
|
@classmethod
|
|
def get_map(cls, size):
|
|
addr = windows.winproxy.VirtualAlloc(0, size, 0x1000, 0x40)
|
|
new_map = (ctypes.c_char * size).from_address(addr)
|
|
new_map.addr = addr
|
|
if new_map.addr == 0:
|
|
raise ctypes.WinError()
|
|
return new_map
|
|
|
|
|
|
class UnixMyMap(MyMap):
|
|
@classmethod
|
|
def get_map(cls, size):
|
|
prot = mmap.PROT_EXEC | mmap.PROT_WRITE | mmap.PROT_READ
|
|
return cls(-1, size, prot=prot)
|
|
|
|
|
|
class CustomAllocator(object):
|
|
int_size = {'32bit': 4, '64bit': 8}
|
|
|
|
def __init__(self):
|
|
self.maps = []
|
|
self.get_new_page(0x1000)
|
|
self.names = []
|
|
|
|
@classmethod
|
|
def get_int_size(cls):
|
|
bits = platform.architecture()[0]
|
|
if bits not in cls.int_size:
|
|
raise ValueError("Unknow platform bits <{0}>".format(bits))
|
|
return cls.int_size[bits]
|
|
|
|
def get_new_page(self, size):
|
|
self.maps.append(MyMap.get_map(size))
|
|
self.cur_offset = 0
|
|
self.cur_page_size = size
|
|
|
|
def reserve_size(self, size):
|
|
if size + self.cur_offset > self.cur_page_size:
|
|
self.get_new_page((size + 0x1000) & ~0xfff)
|
|
addr = self.maps[-1].addr + self.cur_offset
|
|
self.cur_offset += size
|
|
return addr
|
|
|
|
def reserve_int(self, nb_int=1):
|
|
int_size = self.get_int_size()
|
|
return self.reserve_size(int_size * nb_int)
|
|
|
|
def write_code(self, code):
|
|
size = len(code)
|
|
if size + self.cur_offset > self.cur_page_size:
|
|
self.get_new_page((size + 0x1000) & ~0xfff)
|
|
self.maps[-1][self.cur_offset: self.cur_offset + size] = code
|
|
addr = self.maps[-1].addr + self.cur_offset
|
|
self.cur_offset += size
|
|
return addr
|
|
|
|
allocator = CustomAllocator()
|
|
|
|
|
|
def create_function(code, types):
|
|
"""Create a python function that call raw machine code
|
|
|
|
:param str code: Raw machine code that will be called
|
|
:param list types: Return type and parameters type (see :mod:`ctypes`)
|
|
:return: the created function
|
|
:rtype: function
|
|
"""
|
|
func_type = ctypes.CFUNCTYPE(*types)
|
|
addr = allocator.write_code(code)
|
|
res = func_type(addr)
|
|
res.code_addr = addr
|
|
return res
|