hakril-PythonForWindows/tests/test_simple_x86.py

try:
    import capstone
except ImportError as e:
    capstone = None

import pytest

import windows.native_exec.simple_x86 as x86
from windows.native_exec.simple_x86 import *
del Test # Prevent pytest warning

from windows.pycompat import int_types

VERBOSE = False

if capstone:
    disassembleur = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_32)
    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, immediat_accepted=None, expected_result=None, debug=False):
        self.instr_to_test = instr_to_test
        self.expected_result = expected_result
        self.immediat_accepted = immediat_accepted
        self.debug = debug

    def __call__(self, *args):
        if self.debug:
            import pdb;pdb.set_trace()
            pdb.DONE = True
        res = bytes(self.instr_to_test(*args).get_code())
        capres_list = disas(res)
        if len(capres_list) != 1:
            raise AssertionError("Trying to disas an instruction resulted in multiple disassembled instrs")
        capres = capres_list[0]
        print("{0} {1}".format(capres.mnemonic, capres.op_str))
        if len(res) != len(capres.bytes):
            raise AssertionError("Not all bytes have been used by the disassembler")

        if VERBOSE:
            print("  * [CODE] {0}".format(repr(res)))
            print("  * [DISAS] {0} {1}".format(capres.mnemonic, capres.op_str))

        if self.expected_result is not None:
            if "{0} {1}".format(capres.mnemonic, capres.op_str) == 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)))
        self.compare_mnemo(capres)
        self.compare_args(args, capres)

    def compare_mnemo(self, capres):
        expected = self.instr_to_test.__name__.lower()
        if expected != str(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)
        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):
            if isinstance(op_args, str):  # Register
                if cap_op.type != capstone.x86.X86_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))
            elif isinstance(op_args, mem_access):
                self.compare_mem_access(op_args, capres, cap_op)
            else:
                raise ValueError("Unknow argument {0} of type {1}".format(op_args, type(op_args)))

    def compare_mem_access(self, memaccess, capres, cap_op):
        if cap_op.type != capstone.x86.X86_OP_MEM:
            raise AssertionError("Expected Memaccess <{0}> got {1}".format(memaccess, cap_op))
        if memaccess.prefix is not None and capres.prefix[1] != x86_segment_selectors[memaccess.prefix].PREFIX_VALUE:
            try:
                get_prefix = [n for n, x in x86_segment_selectors.items() if x.PREFIX_VALUE == capres.prefix[1]][0]
            except IndexError:
                get_prefix = None
            raise AssertionError("Expected Segment overide <{0}> got {1}".format(memaccess.prefix, get_prefix))
        cap_mem = cap_op.mem
        if memaccess.base is None and cap_mem.base != capstone.x86.X86_REG_INVALID:
            raise AssertionError("Unexpected memaccess base <{0}>".format(capres.reg_name(cap_mem.base)))
        if memaccess.base is not None and capres.reg_name(cap_mem.base) != memaccess.base.lower():
            raise AssertionError("Expected mem.base {0} got {1}".format(memaccess.base.lower(), capres.reg_name(cap_mem.base)))
        if memaccess.index is None and cap_mem.index != capstone.x86.X86_REG_INVALID:
            raise AssertionError("Unexpected memaccess index <{0}>".format(capres.reg_name(cap_mem.base)))
        if memaccess.index is not None and capres.reg_name(cap_mem.index) != memaccess.index.lower():
            raise AssertionError("Expected mem.index {0} got {1}".format(memaccess.index.lower(), capres.reg_name(cap_mem.index)))
        if memaccess.scale != cap_mem.scale and not (memaccess.scale is None and cap_mem.scale == 1):
            raise AssertionError("Expected mem.scale {0} got {1}".format(memaccess.scale, cap_mem.scale))
        if memaccess.disp & 0xffffffff != cap_mem.disp & 0xffffffff:
            raise AssertionError("Expected mem.disp {0} got {1}".format(memaccess.disp, cap_mem.disp))


def test_assembler():
    CheckInstr(Mov)('EAX', 'CR3')
    CheckInstr(Mov)('EDX', 'CR0')
    CheckInstr(Mov)('EDI', 'CR7')

    CheckInstr(Mov)('CR3', 'EAX')
    CheckInstr(Mov)('CR0', 'EDX')
    CheckInstr(Mov)('CR7', 'EDI')

    # Registers
    CheckInstr(Pushad, expected_result="pushal ")()
    CheckInstr(Pushfd)()
    CheckInstr(Popad, expected_result="popal ")()
    CheckInstr(Popfd)()

    CheckInstr(Mov)('EAX', 'ESP')
    CheckInstr(Mov)('ECX', mem('[EAX]'))
    CheckInstr(Mov)('EDX', mem('[ECX + 0x10]'))
    CheckInstr(Mov)('EDX', mem('[EDI * 8 + 0xffff]'))
    CheckInstr(Mov)('EDX', mem('[0x11223344]'))
    CheckInstr(Mov)('EDX', mem('[ESP + EBP * 2 + 0x223344]'))
    CheckInstr(Mov)(mem('[EBP + EBP * 2 + 0x223344]'), 'ESP')
    CheckInstr(Mov)('ESI', mem('[ESI + EDI * 1]'))
    CheckInstr(Mov)('EAX', mem('fs:[0x30]'))
    CheckInstr(Mov)('EDI', mem('gs:[EAX + ECX * 4]'))
    CheckInstr(Mov)('AX', 'AX')
    CheckInstr(Mov)('SI', 'DI')
    CheckInstr(Mov)('AX', 'AX')
    CheckInstr(Mov)('AX', mem('fs:[0x30]'))
    CheckInstr(Mov)('AX', mem('fs:[EAX + 0x30]'))
    CheckInstr(Mov)('AX', mem('fs:[EAX + ECX * 4+0x30]'))
    # Segment selector
    CheckInstr(Mov)('SS', 'ECX')
    CheckInstr(Mov)('ECX', 'SS')
    CheckInstr(Mov)('EDX', 'es')
    CheckInstr(Mov)('EDX', 'cs')
    CheckInstr(Mov)('EDX', 'ds')
    CheckInstr(Mov)('EDX', 'fs')
    CheckInstr(Mov)('fs', 'eax')
    CheckInstr(Mov)('fs', 'eax')


    CheckInstr(Add)('EAX', 8)
    CheckInstr(Add)('EAX', 0xffffffff)
    CheckInstr(Add)("ECX", mem("[EAX + 0xff]"))
    CheckInstr(Add)("ECX", mem("[EAX + 0xffffffff]"))

    CheckInstr(Add)(mem('[EAX]'), 10)
    CheckInstr(Mov)('EAX', mem('fs:[0xfffc]'))
    CheckInstr(Mov)(mem('fs:[0xfffc]'), 0)

    CheckInstr(Push)('ECX')
    CheckInstr(Push)(mem('[ECX + 8]'))

    CheckInstr(Sub)('ECX', 'ESP')
    CheckInstr(Sub)('ECX', mem('[ESP]'))

    CheckInstr(Inc)('EAX')
    CheckInstr(Inc)(mem('[0x42424242]'))
    CheckInstr(Lea)('EAX', mem('[EAX + 1]'))
    CheckInstr(Lea)('ECX', mem('[EDI + -0xff]'))
    CheckInstr(Call)('EAX')
    CheckInstr(Call)(mem('[EAX + ECX * 8]'))
    CheckInstr(Cpuid)()
    CheckInstr(Movsb, expected_result='movsb byte ptr es:[edi], byte ptr [esi]')()
    CheckInstr(Movsd, expected_result='movsd dword ptr es:[edi], dword ptr [esi]')()
    CheckInstr(Xchg)('EAX', 'ESP')

    CheckInstr(Rol)('EAX', 7)
    CheckInstr(Rol)('ECX', 0)

    CheckInstr(Ror)('ECX', 0)
    CheckInstr(Ror)('EDI', 7)
    CheckInstr(Ror)('EDI', -128)

    CheckInstr(Cmp, immediat_accepted=0xffffffff)('EAX', -1)
    CheckInstr(Cmp)('EAX', 0xffffffff)



    CheckInstr(And)('ECX', 'EBX')
    CheckInstr(And)('EAX', 0x11223344)
    CheckInstr(And)('EAX', mem('[EAX + 1]'))
    CheckInstr(And)(mem('[EAX + EAX]'), 'EDX')

    CheckInstr(Or)('ECX', 'EBX')
    CheckInstr(Or)('EAX', 0x11223344)
    CheckInstr(Or)('EAX', mem('[EAX + 1]'))
    CheckInstr(Or)(mem('[EAX + EAX]'), 'EDX')

    CheckInstr(Shr)('EAX', 8)
    CheckInstr(Shr)('EDX', 0x12)
    CheckInstr(Shl)('EAX', 8)
    CheckInstr(Shl)('EDX', 0x12)

    CheckInstr(Not)('EAX')
    CheckInstr(Not)(mem('[EAX]'))

    CheckInstr(Int3)()
    CheckInstr(Int)(0)
    CheckInstr(Int)(3)
    CheckInstr(Int)(0xff)

    CheckInstr(ScasB, expected_result="scasb al, byte ptr es:[edi]")()
    CheckInstr(ScasW, expected_result="scasw ax, word ptr es:[edi]")()
    CheckInstr(ScasD, expected_result="scasd eax, dword ptr es:[edi]")()

    CheckInstr(CmpsB, expected_result="cmpsb byte ptr [esi], byte ptr es:[edi]")()
    CheckInstr(CmpsW, expected_result="cmpsw word ptr [esi], word ptr es:[edi]")()
    CheckInstr(CmpsD, expected_result="cmpsd dword ptr [esi], dword ptr es:[edi]")()


    CheckInstr(x86.Test)('EAX', 'EAX')
    CheckInstr(x86.Test, expected_result="test edi, ecx")('ECX', 'EDI')
    CheckInstr(x86.Test)(mem('[ECX + 0x100]'), 'ECX')
    CheckInstr(x86.Test)('EAX', 0x11223344)
    CheckInstr(x86.Test, immediat_accepted=-1)('EAX', 0xffffffff)
    CheckInstr(x86.Test)('ECX', 0x42)

    CheckInstr(Jmp)('EAX')
    CheckInstr(Jmp)('EDX')
    CheckInstr(Jmp)('EDI')
    CheckInstr(Jmp)(mem('[EAX]'))
    CheckInstr(Jmp)(mem('[EAX + 2]'))
    CheckInstr(Jmp)(mem('[0x12345678]'))

    assert x86.Test(mem('[ECX + 0x100]'), 'ECX').get_code() == x86.Test('ECX', mem('[ECX + 0x100]')).get_code()
    assert Xchg('EAX', 'ECX').get_code() == Xchg('ECX', 'EAX').get_code()

    code = MultipleInstr()
    code += Nop()
    code += Rep + Nop()
    code += Ret()
    print(repr(code.get_code()))
    assert code.get_code() == b"\x90\xf3\x90\xc3"

def test_simple_x64_raw_instruction():
    # Test the fake instruction "raw"
    # 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()
    res += "ret; ret; label :offset_3; ret"
    res += x86.Nop()
    res += x86.Label(":offset_5")
    assert res.get_code() == b"\x90\xc3\xc3\xc3\x90"
    assert res.labels == {":offset_3": 3, ":offset_5": 5}

def test_x86_instr_multiply():
    res = x86.MultipleInstr()
    res += (x86.Nop() * 5)
    res += x86.Ret()
    assert res.get_code() == b"\x90\x90\x90\x90\x90\xc3"

import threading
threads_error = []

def test_x86_multithread_target():
    try:

        assert x86.Mov("ECX", "SS").get_code() == b"\x8c\xd1"
        assert x86.Mov("SS", "ECX").get_code() == b"\x8e\xd1"
        assert x86.Ret().get_code() == b"\xc3"

        res = x86.MultipleInstr()
        res += x86.Mov("ECX", "SS")
        res += x86.Mov("SS", "ECX")
        res += x86.Ret()
        assert res.get_code() == b"\x8c\xd1\x8e\xd1\xc3"
    except Exception as e:
        threads_error.append(e)
        raise
    return True

def test_x86_multithread():
    all_threads = []
    for tnb in range(10):
        t = threading.Thread(target=test_x86_multithread_target)
        all_threads.append(t)

    # import pdb; pdb.set_trace()
    for t in all_threads:
        t.start()
    for t in all_threads:
        t.join()
    assert not threads_error, "syswow call inconsistent with MultiThreading inconsistent"


if capstone is None:
    test_assembler = pytest.mark.skip("Capstone not installed")(test_assembler)

if __name__ == "__main__":
    test_assembler()