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