Merge pull request #54 from hakril/improve_simple_x64

Implem x64 mem_access [RIP+DISP32]

Author: hakril

tests/test_simple_x64.py

@@ -258,6 +258,23 @@ def test_assembler():
     CheckInstr(Not)('RAX')
     CheckInstr(Not)(mem('[RAX]'))
 
+    # Check RIP relative instructions
+    CheckInstr(Mov)(mem('[rip + 0x41414141]'), 0x42424242)
+    CheckInstr(Mov)(mem('[rip + 0x41414141]'), 'R8')
+    CheckInstr(Mov)("RAX", mem('[rip + 0x2]'))
+    CheckInstr(Mov)("RAX", mem('[rip + -5]'))
+    CheckInstr(Mov)("RAX", mem('[rip + 0x41414141]'))
+    CheckInstr(Lea)("RAX", mem('[rip + 0x2]'))
+    CheckInstr(Lea)("RAX", mem('[rip + -5]'))
+    CheckInstr(Lea)("EAX", mem('[rip + 0x2]'))
+    CheckInstr(Lea)("R15", mem('[rip + 0x2]'))
+
+    with pytest.raises(ValueError):
+        CheckInstr(Mov, must_fail=True)('RAX', mem('[RIP + RCX]')) # Invalid mem parsing
+    with pytest.raises(ValueError):
+        CheckInstr(Mov, must_fail=True)('REX', mem('[RIP * 2]')) # Invalid mem parsing
+    with pytest.raises(ValueError):
+        CheckInstr(Mov, must_fail=True)('REX', mem('[RCX + RIP]')) # Invalid mem parsing
 
     CheckInstr(ScasB, expected_result="scasb al, byte ptr [rdi]")()
     CheckInstr(ScasW, expected_result="scasw ax, word ptr [rdi]")()
@@ -297,6 +314,10 @@ def test_simple_x64_raw_instruction():
     # By emetting a multi-char nop manually
     CheckInstr(Raw, expected_result="nop word ptr [rax + rax]")("66 0F 1F 84 00 00 00 00 00")
 
+def test_simple_x64_assemble_raw_one_byte():
+    # Test the fake instruction "raw" inside an assemble that may translate str to int
+    assert x64.assemble("ret; raw 90; ret") ==  b"\xc3\x90\xc3"
+
 def test_x64_32b_register_lower_upper():
     assert x64.assemble("mov eax, 42") == x64.assemble("mov EAX, 42")
     assert x64.Mov("eax", 42).get_code() == x64.Mov("EAX", 42).get_code()

tests/test_simple_x86.py

@@ -246,6 +246,10 @@ def test_simple_x64_raw_instruction():
     # By emetting a multi-char nop manually
     CheckInstr(Raw, expected_result="nop word ptr [eax + eax]")("66 0F 1F 84 00 00 00 00 00")
 
+def test_simple_x86_assemble_raw_one_byte():
+    # Test the fake instruction "raw" inside an assemble that may translate str to int
+    assert x86.assemble("ret; raw 90; ret") ==  b"\xc3\x90\xc3"
+
 def test_x86_multiple_instr_add_instr_and_str():
     res = x86.MultipleInstr()
     res += x86.Nop()

windows/native_exec/simple_x64.py

@@ -136,6 +136,7 @@ def create_prefix(name, value):
 reg_order = ['RAX', 'RCX', 'RDX', 'RBX', 'RSP', 'RBP', 'RSI', 'RDI']
 new_reg_order = ['R8', 'R9', 'R10', 'R11', 'R12', 'R13', 'R14', 'R15']
 x64_regs = reg_order + new_reg_order
+# Rip is now addressable but is a special case
 
 registers_32_bits = {'R15D': 'R15', 'R14D': 'R14', 'ESP': 'RSP', 'R9D': 'R9',
 'EDI': 'RDI', 'R11D': 'R11', 'R8D': 'R8', 'R10D': 'R10', 'EAX': 'RAX',
@@ -170,6 +171,10 @@ def is_32b_reg(name):
         except AttributeError:  # Not a string
             return False
 
+    @staticmethod
+    def is_rip(name):
+        return name is not None and name.upper() == "RIP"
+
     @staticmethod
     def is_mem_acces(data):
         return isinstance(data, mem_access)
@@ -201,6 +206,9 @@ def create_displacement(base=None, index=None, scale=None, disp=0, prefix=None):
         raise ValueError("Cannot create displacement with scale and no index")
     if scale and index.upper() == "RSP":
         raise ValueError("Cannot create displacement with index == RSP")
+    if base and X64.is_rip(base):
+        if index is not None or scale is not None:
+            raise ValueError("Bad use of RIP in mem expression <{0}> : only [RIP+DISP32] is possible".format(data))
     return mem_access(base, index, scale, disp, prefix)
 
 
@@ -251,10 +259,14 @@ def mem(data):
             parsed_items['index'] = index
         else:
             # displacement / base / index alone
-            if X64.is_reg(item):
+            if X64.is_reg(item) or X64.is_rip(item):
                 if 'base' not in parsed_items:
                     parsed_items['base'] = item
                     continue
+                # RIP is only usable on the format [RIP+DISP32]
+                # So no index when base is RIP & no index as RIP
+                if X64.is_rip(item) or X64.is_rip(parsed_items['base']):
+                    raise ValueError("Bad use of RIP in mem expression <{0}> : only [RIP+DISP32] is possible".format(data))
                 # Already have base + index -> cannot avec another register in expression
                 if 'index' in parsed_items:
                     raise ValueError("Multiple index / index*scale in mem expression <{0}>".format(data))
@@ -656,18 +668,18 @@ def __init__(self, arg1, arg2, reversed, instr_state):
             instr_state.prefixes.append(x64_segment_selectors[arg2.prefix])
         # # ARG1 : REG
         # # ARG2 : [MEM]
-        # # this encode [rip + disp]
-        # # TODO :)
-        # if X64.mem_access_has_only(arg2, ["disp"]):
-        #     self.mod = BitArray(2, "00")
-        #     self.setup_reg_as_register(arg1)
-        #     self.rm = BitArray(3, "101")
-        #     try:
-        #         self.after = BitArray.from_string(accept_as_32immediat(arg2.disp))
-        #     except ImmediatOverflow:
-        #         raise ImmediatOverflow("Interger32 overflow for displacement {0}".format(hex(arg2.disp)))
-        #     self.direction = not reversed
-        #     return
+
+        # Special case [RIP+DISP32]
+        if X64.is_rip(arg2.base):
+            self.mod = BitArray(2, "00")
+            self.setup_reg_as_register(arg1)
+            self.rm = BitArray(3, "101")
+            try:
+                self.after = BitArray.from_string(accept_as_32immediat(arg2.disp))
+            except ImmediatOverflow:
+                raise ImmediatOverflow("Interger32 overflow for displacement {0}".format(hex(arg2.disp)))
+            self.direction = not reversed
+            return
 
         # Those registers cannot be addressed without SIB
         FIRE_UP_SIB = not arg2.base or arg2.base.upper() in ["RSP", "RBP"] or arg2.index
@@ -1098,6 +1110,8 @@ class Retf(Instruction):
     default_32_bits = True
     encoding = [(RawBits.from_int(8, 0xcb),)]
 
+class Retfq(Instruction): # Give an explicit name to retf-64bits as it cannot be guessed from the mnemonic alone
+    encoding = [(RawBits.from_int(8, 0x48), RawBits.from_int(8, 0xcb))]
 
 class Retf32(Instruction):
     encoding = [(RawBits.from_int(8, 0xcb),)]
@@ -1115,11 +1129,15 @@ def JmpAt(addr):
 
 class Raw(Instruction):
     """Output raw data"""
+
     def __init__(self, *initial_args):
         if len(initial_args) != 1:
             raise ValueError("raw 'opcode' only accept one argument")
-        # Accept space
-        self.data = binascii.unhexlify(initial_args[0].replace(" ", ""))
+        if isinstance(initial_args[0], int):
+            self.data = struct.pack("B", initial_args[0])
+        else:
+            # Accept space
+            self.data = binascii.unhexlify(initial_args[0].replace(" ", ""))
 
     def get_code(self):
         return self.data
@@ -1300,6 +1318,12 @@ def assemble_instructions_generator(str):
         except:
             raise ValueError("Unknow mnemonic <{0}>".format(mnemo))
 
+        if issubclass(instr_object, Raw):
+            # Raw should received the raw buffer as it expect encoded hex
+            # The transformation may transform 'raw 9090' (nopnop) as 0n9090
+            # If other fake-instr need this : make a class attribute
+            yield instr_object(*args_raw)
+            continue
         args = []
         if args_raw:
             for arg in args_raw[0].split(","):

windows/native_exec/simple_x86.py

@@ -1128,6 +1128,13 @@ def assemble_instructions_generator(str):
         except:
             raise ValueError("Unknow mnemonic <{0}>".format(mnemo))
 
+        if issubclass(instr_object, Raw):
+            # Raw should received the raw buffer as it expect encoded hex
+            # The transformation may transform 'raw 9090' (nopnop) as 0n9090
+            # If other fake-instr need this : make a class attribute
+            yield instr_object(*args_raw)
+            continue
+
         args = []
         if args_raw:
             for arg in args_raw[0].split(","):