43 lines
1.1 KiB
Python
43 lines
1.1 KiB
Python
from enum import IntEnum
|
|
|
|
import lief
|
|
|
|
import cfg_parser
|
|
from utils import Instructions, hex_address_to_memory_representation, is_32b, is_little_endian
|
|
|
|
|
|
class Registers(IntEnum):
|
|
EAX = 0xC0
|
|
EBX = 0xC3
|
|
ECX = 0xC1
|
|
EDX = 0xC2
|
|
ESI = 0xC6
|
|
EDI = 0xC7
|
|
EBP = 0xC5
|
|
# ESP = 0xC4
|
|
|
|
|
|
def generate_reg_init_code(cfg, pe: lief.PE.Binary, wave: int, wave_entry: int) -> list[int]:
|
|
code = []
|
|
reg_values = cfg_parser.parse_bb_registers(cfg, wave, 0)
|
|
for reg in reg_values:
|
|
if reg not in Registers.__members__:
|
|
continue
|
|
new_instruction = (
|
|
Instructions.MOV_REG
|
|
+ [Registers[reg]]
|
|
+ hex_address_to_memory_representation(
|
|
reg_values[reg].strip(),
|
|
is_32b(pe),
|
|
is_little_endian(pe),
|
|
)
|
|
)
|
|
for byte in new_instruction:
|
|
code.append(byte)
|
|
|
|
# add ret to actual OEP
|
|
code += Instructions.PUSH + hex_address_to_memory_representation(hex(wave_entry), is_32b(pe), is_little_endian(pe)) # push addr
|
|
|
|
code += Instructions.RET
|
|
|
|
return code
|