From 1bc7818552b47670da9a61d671c5c1af7f1660e9 Mon Sep 17 00:00:00 2001 From: hakril Date: Sat, 25 Jan 2025 00:50:25 +0100 Subject: [PATCH] Still playing with simple_arm64 --- tests/test_simple_arm64.py | 113 +++++++++++++++++++------- windows/native_exec/simple_arm64.py | 122 +++++++++++++++++++++++++--- 2 files changed, 194 insertions(+), 41 deletions(-) diff --git a/tests/test_simple_arm64.py b/tests/test_simple_arm64.py index cab9de4..ee5fa19 100644 --- a/tests/test_simple_arm64.py +++ b/tests/test_simple_arm64.py @@ -35,8 +35,22 @@ class CheckInstr(object): self.expected_result = expected_result self.must_fail = must_fail self.debug = debug + self.callargs = None + def __call__(self, *args): + assert args is not None + self.callargs = args + return self + + def __repr__(self): + if self.must_fail: + return "MustFail:{0}{1}".format(self.instr_to_test.__name__, self.callargs) + return "{0}{1}".format(self.instr_to_test.__name__, self.callargs) + + def dotest(self): + assert self.callargs is not None + args = self.callargs try: if self.debug: import pdb;pdb.set_trace() @@ -70,6 +84,7 @@ class CheckInstr(object): raise AssertionError("Not all bytes have been used by the disassembler") self.compare_mnemo(capres) self.compare_args(args, capres) + return True def compare_mnemo(self, capres): expected = self.instr_to_test.__name__.lower() @@ -80,9 +95,15 @@ class CheckInstr(object): def compare_args(self, args, capres): capres_op = list(capres.operands) - if len(args) != len(capres_op): - raise AssertionError("Expected {0} operands got {1}".format(len(args), len(capres_op))) - for op_args, cap_op in zip(args, capres_op): + # We may have != number of operand as shift are: + # - arguments for simple_arm64 + # - atribute of immediat for capstone + if not len(capres_op) <= len(args): + raise AssertionError("Expected at most {0} operands got {1}".format(len(args), len(capres_op))) + + opargit = iter(args) # allow manually using next() to get next simple_arm64 arg for shift compare + # capres_op must be first in zip (as its smaller) or last next(opargit) will be consommed by zip + for cap_op, op_args in zip(capres_op, opargit): if isinstance(op_args, str): # Register if cap_op.type != capstone.arm64.ARM64_OP_REG: raise AssertionError("Expected args {0} operands got {1}".format(op_args, capres_op)) @@ -91,9 +112,39 @@ class CheckInstr(object): elif isinstance(op_args, int_types): if (op_args != cap_op.imm) and not (self.immediat_accepted and self.immediat_accepted == cap_op.imm): raise AssertionError("Expected Immediat <{0}> got {1}".format(op_args, cap_op.imm)) + cap_shift = cap_op.shift + if not (cap_shift.type == cap_shift.value == 0): + self.compare_shift(next(opargit), cap_shift) else: raise ValueError("Unknow argument {0} of type {1}".format(op_args, type(op_args))) + # Check that no argument were unused in args + # As args + shift should perfectly match the capres_op + sentinel = object() + nextarg = next(opargit, sentinel) + if nextarg != sentinel: + # Ignore a leading LSL #0 shift, as it should be authorized but not displayed by disassembler + shift = Shift.parse(nextarg) + if not (shift.type == "LSL" and shift.value == 0): + raise ValueError("Non consomated argument: {0} (probable non-encoded shift)".format(nextarg)) + + SHIFT_TYPE_TO_CAPSTONE = { + "LSL": capstone.arm64.ARM64_SFT_LSL, + "LSR": capstone.arm64.ARM64_SFT_LSR, + "ASR": capstone.arm64.ARM64_SFT_ASR, + "ROR": capstone.arm64.ARM64_SFT_ROR, + # "MSL": apstone.arm64.ARM64_SFT_MSL # Not yet used in PFW + } + + def compare_shift(self, shiftstr, cap_shift): + shift = Shift.parse(shiftstr) + if not self.SHIFT_TYPE_TO_CAPSTONE[shift.type] == cap_shift.type: + raise ValueError("Shift type mismatch: expected {0} got {1}".format(shift.type, cap_shift.type)) + if not shift.value == cap_shift.value: + raise ValueError("Shift value mismatch: expected {0} got {1}".format(shift.value, cap_shift.value)) + return True + + def test_shift_parsing(): assert Shift.parse("LSL #0") assert Shift.parse("LSL #12") @@ -110,32 +161,36 @@ def test_shift_parsing(): assert not Shift.parse("LSX ##1") assert not Shift.parse("LSX #") -def test_assembler(): - CheckInstr(Add)('W0', 'W0', 0) - CheckInstr(Add)('W1', 'W0', 0) - CheckInstr(Add)('W30', 'W12', 0) - CheckInstr(Add)('W0', 'W0', 1) +@pytest.mark.parametrize("checkinstr", [ + CheckInstr(Add)('W0', 'W0', 0), + CheckInstr(Add)('W1', 'W0', 0), + CheckInstr(Add)('W30', 'W12', 0), + CheckInstr(Add)('W0', 'W0', 1), + CheckInstr(Add)('X0', 'X0', 0), + CheckInstr(Add)('X30', 'X12', 0), + CheckInstr(Add)('X0', 'X0', 1), + CheckInstr(Add)('X11', 'X12', 0x123), + CheckInstr(Add)('X11', 'X12', 0x123, "LSL #0"), + CheckInstr(Add)('X11', 'X12', 0x123, "LSL #12"), + CheckInstr(Add, must_fail=True)('X11', 'W12', 0x123), # Bitness mismatch + CheckInstr(Add, must_fail=True)('BADREG', 'X12', 0), + CheckInstr(Add, must_fail=True)('X11', 'X12', 0x123, "LSL #1234"), + CheckInstr(Add, must_fail=True)('X11', 'X12', 0x12345678), - CheckInstr(Add)('X0', 'X0', 0) - CheckInstr(Add)('X30', 'X12', 0) - CheckInstr(Add)('X0', 'X0', 1) - CheckInstr(Add)('X11', 'X12', 0x123) - # CheckInstr(Add)('X11', 'X12', 0x123, "LSL #0") - CheckInstr(Add)('X11', 'X12', 0x123, "LSL #12") + CheckInstr(Movz)('X0', 0), + CheckInstr(Movz)('X0', 0, "LSL #32"), + CheckInstr(Movz)('X18', 0, "LSL #48"), + CheckInstr(Movz)('W18', 0, "LSL #16"), + CheckInstr(Movz, must_fail=True)('X0', 0, "LSL #12"), # Invalid LSL for MovWideImmediat + CheckInstr(Movz, must_fail=True)('W0', 0, "LSL #32"), + CheckInstr(Movz, must_fail=True)('X0', 0, "ROR #32"), - # Error test todo - # CheckInstr(Add)('X11', 'W12', 0x123) - with pytest.raises(ValueError): - CheckInstr(Add)('BADREG', 'X12', 0) - with pytest.raises(ValueError): - CheckInstr(Add)('X11', 'X12', 0x123, "LSL #1234") + CheckInstr(Movk)('X0', 0x1234, "LSL #32"), + CheckInstr(Movk)('X18', 0x5678, "LSL #48"), - with pytest.raises(ValueError): - # Immediat too big for encoding - CheckInstr(Add)('X11', 'X12', 0x12345678) - - CheckInstr(Ret)("X0") - CheckInstr(Ret, expected_result="ret ")("X30") - CheckInstr(Ret)() - with pytest.raises(ValueError): - CheckInstr(Ret)("W0") + CheckInstr(Ret)("X0"), + CheckInstr(Ret, expected_result="ret ")("X30"), + CheckInstr(Ret)(), +], ids=CheckInstr.__repr__) +def test_instruction_assembling(checkinstr): + assert checkinstr.dotest() \ No newline at end of file diff --git a/windows/native_exec/simple_arm64.py b/windows/native_exec/simple_arm64.py index 6b53634..ea3803c 100644 --- a/windows/native_exec/simple_arm64.py +++ b/windows/native_exec/simple_arm64.py @@ -113,7 +113,7 @@ class InstructionEncoding(object): return (accept_sp and (arg in [SP, WSP])) or arg in ALL_REGISTER @classmethod - def is_imm12(self, arg): + def is_imm(self, arg): try: value = int(arg) except (ValueError, TypeError): @@ -127,9 +127,9 @@ class InstructionEncoding(object): @classmethod def gen(cls, **encoding_array): - class GeneratedEncoding(cls): + class GeneratedEncodingCls(cls): ENCODING_VALUES = encoding_array - return GeneratedEncoding + return GeneratedEncodingCls # Instruction filing at instanciation @@ -214,7 +214,6 @@ class AddSubtractImmediate(DataProcessingImmediate): if shift not in [("LSL", 0), ("LSL", 12)]: raise ValueError("Invalid shift for instruction: {0}".format(shift)) if shift == ("LSL", 12): - import pdb;pdb.set_trace() self.sh[:] = bytearray((1,)) @@ -222,10 +221,45 @@ class AddSubtractImmediate(DataProcessingImmediate): def accept_arg(cls, argsdict): return (cls.is_register(argsdict[0], accept_sp=True) and cls.is_register(argsdict[1], accept_sp=True) and - cls.is_imm12(argsdict[2]) and + cls.is_imm(argsdict[2]) and cls.is_shift(argsdict.get(3))) +# C4.1.93.6 Logical (immediate) +# Wtf : https://kddnewton.com/2022/08/11/aarch64-bitmask-immediates.html + +class DataProcessingLogicalImmediate(DataProcessingImmediate): + def __init__(self, argsdict): + super(DataProcessingLogicalImmediate, self).__init__() + self.sf = self.bits[31:32] + self.opc = self.bits[29:31] + self.bits[23:29] = bytearray(reversed((1, 0, 0, 1, 0, 0))) + self.N = self.bits[22:23] + self.immr = self.bits[16:22] + self.imms = self.bits[10:16] + self.rn = self.bits[5:10] + self.rd = self.bits[0:5] + + self.setup_fixed_values() + # Change instruction based of parameter + self.setup_register(self.rd, argsdict[0]) + self.setup_register(self.rn, argsdict[1]) + self.setup_bitmask_imm(self.imm12, argsdict[2]) + + @classmethod + def accept_arg(cls, argsdict): + return (cls.is_register(argsdict[0], accept_sp=True) and + cls.is_register(argsdict[1], accept_sp=True) and + cls.is_bitmask_imm(argsdict[2])) + + @classmethod + def is_bitmask_imm(*args, **kwargs): + raise NotImplementedError("is_bitmask_imm") + + def setup_bitmask_imm(*args, **kwargs): + raise NotImplementedError("setup_bitmask_imm") + + class MovWideImmediat(DataProcessingImmediate): def __init__(self, argsdict): super(MovWideImmediat, self).__init__() @@ -242,13 +276,23 @@ class MovWideImmediat(DataProcessingImmediate): self.setup_register(self.rd, argsdict[0]) self.setup_immediat(self.imm16, argsdict[1]) - assert argsdict.get(3) is None, "SHIFT NOT IMPLEMENTED YET" + shift = Shift.parse(argsdict.get(2)) + if not shift: + return + if shift.type != "LSL": + raise ValueError("Invalid shift type for {0} : {1}".format(type(self).__name__, shift.value)) + if shift.value not in (0, 16 ,32, 48): + raise ValueError("Invalid shift value for {0} : {1}".format(type(self).__name__, shift.value)) + if self.bitness == 32 and shift.value > 16: + raise ValueError("Invalid shift value for 32bits encoding of {0} : {1}".format(type(self).__name__, shift.value)) + + self.setup_immediat(self.hw, shift.value // 16) @classmethod def accept_arg(cls, argsdict): return (cls.is_register(argsdict[0], accept_sp=True) and - cls.is_imm12(argsdict[1]) and + cls.is_imm(argsdict[1]) and cls.is_shift(argsdict.get(2))) @@ -293,24 +337,71 @@ class RetEncoding(UnconditionalBranchRegister.gen(opc=0b10, op2=0b11111, op3=0, class DataProcessingRegister(InstructionEncoding): def __init__(self): super(DataProcessingRegister, self).__init__() - self.bits[26:29] = bytearray((0,0,1)) self.op0 = self.bits[30:31] self.op1 = self.bits[28:29] self.bits[25:28] = bytearray(reversed((1, 0, 1))) self.op2 = self.bits[21:25] self.op3 = self.bits[10:16] +class DataProcessingLogicalShiftedRegister(DataProcessingRegister): + def __init__(self, argsdict): + super(DataProcessingLogicalShiftedRegister, self).__init__() + self.sf = self.bits[31:32] + self.opc = self.bits[29:31] + self.bits[24:29] = bytearray(reversed((0, 1, 0, 1, 0))) + self.shift = self.bits[22:24] + self.N = self.bits[21:22] + self.rm = self.bits[16:21] + self.imm6 = self.bits[10:16] + self.rn = self.bits[5:10] + self.rd = self.bits[0:5] + + self.setup_fixed_values() + # Change instruction based of parameter + self.setup_register(self.rd, argsdict[0]) + self.setup_register(self.rn, argsdict[1]) + self.setup_register(self.rm, argsdict[2]) + + shift = Shift.parse(argsdict.get(3)) + if not shift: + return + # Is this mapping generic ? Store ir somewhere ? + # Is the shift size logic repeatable and factorisable ? + if self.bitness == 32 and shift.value > 31: + raise ValueError("Invalid shift value for 32bits encoding of {0} : {1}".format(type(self).__name__, shift.value)) + + SHIFT_MAPPING = {"LSL": 0b00, "LSR": 0b01, "ASR": 0b10, "ROR": 0b11} + self.setup_immediat(self.shift, SHIFT_MAPPING[shift.type]) + self.setup_immediat(self.imm6, shift.value) + + + @classmethod + def accept_arg(cls, argsdict): + return (cls.is_register(argsdict[0]) and + cls.is_register(argsdict[1]) and + cls.is_register(argsdict[2]) and + cls.is_shift(argsdict.get(3))) + # An instruction is a Name that can have multiple encoding # It's the class we instanciate to assemble instructions -# Add X0, X0, IMM -# Add X0, X0, X0 +# C6.2.270 ORR (immediate) +# C6.2.271 ORR (shifted register) + +# there also seem to exist "alias instructions" like "mov" +# That just map to others instruction when specific condition are met on the params + class Instruction(object): encoding = [] def __init__(self, *args): argsdict = dict(enumerate(args)) # Like a list but allow arg.get(4) - for encodcls in self.encoding: + for i, encodcls in enumerate(self.encoding): + # Late rewrite of GeneratedEncodingCls classname for better message error + if encodcls.__name__ == "GeneratedEncodingCls": + encodcls.__name__ = "{0}Encoding{1}".format(type(self).__name__, i) + + if encodcls.accept_arg(argsdict): self.encoded = encodcls(argsdict) return @@ -346,10 +437,17 @@ class Ret(Instruction): # C6.2.254 -class MovZ(Instruction): +class Movz(Instruction): encoding = [MovWideImmediat.gen(opc=0b10)] +class Movk(Instruction): + encoding = [MovWideImmediat.gen(opc=0b11)] +# The encoding for "mov reg, reg" :D +# C6.2.271 +# Todo: Instruction like "mov" that dispatch to other instruction encoding based on more precise condition on param ? +class Orr(Instruction): + encoding = [DataProcessingLogicalShiftedRegister.gen(opc=0b01)] class MultipleInstr(object): INSTRUCTION_SIZE = 4