from rekall import addrspace
from rekall import testlib
from rekall.plugins.windows.malware import apihooks


class TestHookHeuristics(testlib.RekallBaseUnitTestCase):
    """Test the hook detection heuristic.

    The actual test cases are generated using the nasm assembler in:

    rekall/src/hooks/amd64.asm and rekall/src/hooks/i386.asm
    """

    def testHook(self):
        session = self.MakeUserSession()

        # The target address should be fixed at this offset.
        target = 0x100

        heuristic = apihooks.HookHeuristic(session=session)

        profile = session.profile = session.LoadProfile("tests/hooks")
        for arch in ["AMD64", "I386"]:
            for test_case in profile.data[arch]:
                offset = test_case["offset"]
                # Test case data is the assembly snippet mapped at the specified
                # offset in the address space.
                address_space = addrspace.BufferAddressSpace(
                    data=test_case["data"].decode("base64"),
                    session=session, base_offset=offset)

                function = session.profile.Function(
                    offset=offset, vm=address_space, name=test_case["name"],
                    mode=arch)

                # Detect the jump in this function
                destination = heuristic.Inspect(function)

                # All hooks in test cases go to the same target offset (0x100).
                self.assertEqual(destination, target)
