import lief
import cfg_parser
from utils import hex_address_to_memory_representation

def generate_reg_init_code(cfg, pe: lief.PE.Binary,wave:int, wave_entry: int) -> list[int]:
    code = []
    # initiate registry values
    reg_to_inst_code = {
        "EAX": 0xC0,
        "EBX": 0xC3,
        "ECX": 0xC1,
        "EDX": 0xC2,
        "ESI": 0xC6,
        "EDI": 0xC7,
        "EBP": 0xC5,
        # "ESP": 0xC4,
    }
    reg_values = cfg_parser.parse_bb_registers(cfg, wave, 0)
    for reg in reg_values:
        if reg not in reg_to_inst_code:
            continue
        new_instruction = [
            0xC7,
            reg_to_inst_code[reg],
        ] + hex_address_to_memory_representation(
            reg_values[reg].strip(),
            pe.abstract.header.is_32,
            pe.abstract.header.endianness == lief.Header.ENDIANNESS.LITTLE,
        )
        for byte in new_instruction:
            code.append(byte)

    # add ret to actual OEP
    code += [0x68] + hex_address_to_memory_representation(
        hex(wave_entry),
        pe.abstract.header.is_32,
        pe.abstract.header.endianness == lief.Header.ENDIANNESS.LITTLE,
    ) # push addr

    code += [0xC3] # ret

    return code