| 12 | |
| 13 | |
| 14 | class CustomAllocator(object): |
| 15 | int_size = {'32bit': 4, '64bit': 8} |
| 16 | |
| 17 | def __init__(self): |
| 18 | self.maps = [] |
| 19 | self.cur_offset = 0 |
| 20 | self.cur_page_size = 0 # Force get_new_page on first request |
| 21 | self.names = [] |
| 22 | |
| 23 | @classmethod |
| 24 | def get_int_size(cls): |
| 25 | bits = platform.architecture()[0] |
| 26 | if bits not in cls.int_size: |
| 27 | raise ValueError("Unknow platform bits <{0}>".format(bits)) |
| 28 | return cls.int_size[bits] |
| 29 | |
| 30 | def get_new_page(self, size): |
| 31 | addr = windows.winproxy.VirtualAlloc(0, size, 0x1000, gdef.PAGE_EXECUTE_READWRITE) |
| 32 | mymap = (ctypes.c_char * size).from_address(addr) |
| 33 | mymap.addr = addr |
| 34 | self.maps.append(mymap) |
| 35 | self.cur_offset = 0 |
| 36 | self.cur_page_size = size |
| 37 | |
| 38 | def reserve_size(self, size): |
| 39 | if size + self.cur_offset > self.cur_page_size: |
| 40 | self.get_new_page((size + 0x1000) & ~0xfff) |
| 41 | addr = self.maps[-1].addr + self.cur_offset |
| 42 | self.cur_offset += size |
| 43 | return addr |
| 44 | |
| 45 | def reserve_int(self, nb_int=1): |
| 46 | int_size = self.get_int_size() |
| 47 | return self.reserve_size(int_size * nb_int) |
| 48 | |
| 49 | def write_code(self, code): |
| 50 | size = len(code) |
| 51 | if size + self.cur_offset > self.cur_page_size: |
| 52 | self.get_new_page((size + 0x1000) & ~0xfff) |
| 53 | self.maps[-1][self.cur_offset: self.cur_offset + size] = code |
| 54 | addr = self.maps[-1].addr + self.cur_offset |
| 55 | self.cur_offset += size |
| 56 | return addr |
| 57 | |
| 58 | def close(self): |
| 59 | maps = self.maps |
| 60 | self.maps = [] |
| 61 | self.cur_offset = 0 |
| 62 | self.cur_page_size = 0 |
| 63 | if getattr(sys, "path", None) is None: |
| 64 | # Path is None -> Python shutdown |
| 65 | return |
| 66 | for mymap in maps: |
| 67 | windows.winproxy.VirtualFree(mymap.addr, dwFreeType=gdef.MEM_RELEASE) |
| 68 | |
| 69 | def __del__(self): |
| 70 | self.close() |
| 71 | |