code formatting and newly created enums to clean up code

reginit.py

@@ -1,42 +1,43 @@
-import lief
-import cfg_parser
-from utils import hex_address_to_memory_representation
+from enum import IntEnum
 
-def generate_reg_init_code(cfg, pe: lief.PE.Binary,wave:int, wave_entry: int) -> list[int]:
+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 = []
-    # 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:
+        if reg not in Registers.__members__:
             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,
+        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 += [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 += Instructions.PUSH + hex_address_to_memory_representation(hex(wave_entry), is_32b(pe), is_little_endian(pe))  # push addr
 
-    code += [0xC3] # ret
+    code += Instructions.RET
 
     return code