mirror of
https://github.com/hakril/PythonForWindows
synced 2026-06-08 14:31:45 +00:00
hakril
211 lines
8.5 KiB
Python
211 lines
8.5 KiB
Python
try:
|
|
import capstone
|
|
except ImportError as e:
|
|
capstone = None
|
|
import pytest
|
|
|
|
import windows.native_exec.simple_arm64 as arm64
|
|
from windows.native_exec.simple_arm64 import *
|
|
|
|
from windows.pycompat import int_types
|
|
|
|
if capstone:
|
|
disassembleur = capstone.Cs(capstone.CS_ARCH_ARM64, capstone.CS_MODE_ARM)
|
|
disassembleur.detail = True
|
|
|
|
@pytest.fixture
|
|
def need_capstone():
|
|
if capstone is None:
|
|
raise pytest.skip("Capstone is not installed")
|
|
return True
|
|
|
|
pytestmark = pytest.mark.usefixtures("need_capstone")
|
|
|
|
|
|
def disas(x):
|
|
return list(disassembleur.disasm(x, 0))
|
|
|
|
|
|
class CheckInstr(object):
|
|
def __init__(self, instr_to_test, expected_result=None, immediat_accepted=None, must_fail=None, debug=False):
|
|
self.instr_to_test = instr_to_test
|
|
self.immediat_accepted = immediat_accepted
|
|
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()
|
|
pdb.DONE = True
|
|
arm64.DEBUG = self.debug
|
|
res = bytes(self.instr_to_test(*args).get_code())
|
|
if self.debug:
|
|
print(repr(res))
|
|
except ValueError as e:
|
|
if self.must_fail == True:
|
|
return True
|
|
else:
|
|
raise
|
|
else:
|
|
if self.must_fail:
|
|
raise ValueError("Instruction did not failed as expected")
|
|
capres_list = disas(res)
|
|
if not capres_list:
|
|
raise AssertionError("Trying to disas an instruction resulted no disassembled instr")
|
|
if len(capres_list) != 1:
|
|
raise AssertionError("Trying to disas an instruction resulted in multiple disassembled instrs: {0}".format(capres_list))
|
|
capres = capres_list[0]
|
|
print("{0} {1}".format(capres.mnemonic, capres.op_str))
|
|
if self.expected_result is not None:
|
|
if "{0} {1}".format(capres.mnemonic, capres.op_str).strip() == self.expected_result:
|
|
return True
|
|
else:
|
|
raise AssertionError("Expected result <{0}> got <{1}>".format(self.expected_result, "{0} {1}".format(capres.mnemonic, capres.op_str)))
|
|
if len(res) != len(capres.bytes):
|
|
print("<{0}> vs <{1}>".format(repr(res), repr(capres.bytes)))
|
|
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()
|
|
if expected != capres.mnemonic:
|
|
raise AssertionError("Expected menmo {0} got {1}".format(expected, str(capres.mnemonic)))
|
|
return True
|
|
|
|
def compare_args(self, args, capres):
|
|
capres_op = list(capres.operands)
|
|
# 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))
|
|
if op_args.lower() != capres.reg_name(cap_op.reg).lower():
|
|
raise AssertionError("Expected register <{0}> got {1}".format(op_args.lower(), capres.reg_name(cap_op.reg).lower()))
|
|
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))
|
|
|
|
if capstone:
|
|
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")
|
|
assert Shift.parse("LSL #1")
|
|
assert Shift.parse("LSR #1")
|
|
|
|
assert Shift.parse("ROR #0").type == "ROR"
|
|
assert Shift.parse("LSL #0").type == "LSL"
|
|
assert Shift.parse("LSL #0").value == 0
|
|
assert Shift.parse("LSL #1").type == "LSL"
|
|
assert Shift.parse("LSL #1").value == 1
|
|
|
|
assert not Shift.parse("LSX #1")
|
|
assert not Shift.parse("LSX ##1")
|
|
assert not Shift.parse("LSX #")
|
|
|
|
@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(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"),
|
|
|
|
|
|
CheckInstr(Orr)('X0', 'X18', 'XZR'),
|
|
# Official encoding for this in ARM64 ref
|
|
CheckInstr(Orr, expected_result="mov x0, x18")('X0', 'XZR', 'X18'),
|
|
CheckInstr(Orr, must_fail=True)('X0', 'WZR', 'X18'), # Size mismatch
|
|
CheckInstr(Orr, must_fail=True)('W0', 'XZR', 'W18'), # Size mismatch
|
|
|
|
CheckInstr(Movk)('X0', 0x1234, "LSL #32"),
|
|
CheckInstr(Movk)('X18', 0x5678, "LSL #48"),
|
|
|
|
CheckInstr(Ret)("X0"),
|
|
CheckInstr(Ret, expected_result="ret")("X30"),
|
|
CheckInstr(Ret)(),
|
|
|
|
# Virtual instruction that dispatch to something else:
|
|
# Ex: "mov reg1, re2" -> "orr reg1, xzr, reg2"
|
|
|
|
CheckInstr(Mov)('X0', 'X18'),
|
|
CheckInstr(Mov)('W0', 'W18'),
|
|
CheckInstr(Mov, must_fail=True)('X0', 'W18'),
|
|
CheckInstr(Mov, must_fail=True)('X0', 'X18', 'X12'),
|
|
|
|
|
|
], ids=CheckInstr.__repr__)
|
|
def test_instruction_assembling(checkinstr):
|
|
assert checkinstr.dotest() |