# Rekall Memory Forensics
#
# Based on original code by:
# Copyright (C) 2007-2013 Volatility Foundation
#
# Authors:
# Michael Hale Ligh <michael.ligh@mnin.org>
#
# This code:
# Copyright 2014 Google Inc. All Rights Reserved.
#
# Authors:
# Michael Cohen <scudette@google.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or (at
# your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
#

from builtins import hex
from builtins import object
import struct

from rekall import testlib

from rekall.plugins.windows import common
from rekall.plugins.overlays.windows import pe_vtypes
from rekall_lib import utils


class DecodingError(Exception):
    """Raised when unable to decode an instruction."""


class HookHeuristic(object):
    """A Hook heuristic detects possible hooks.

    This heuristic emulates some common CPU instructions to try and detect
    control flow jumps within the first few instructions of a function.

    These are essentially guesses based on the most common hook types. Be aware
    that these are pretty easy to defeat which will cause the hook to be missed.

    See rekall/src/hooks/amd64.asm and rekall/src/hooks/i386.asm For the test
    cases which illustrate the type of hooks that we will detect.
    """

    def __init__(self, session=None):
        self.session = session
        self.Reset()

    def Reset(self):
        # Keep track of registers, stack and main memory.
        self.regs = {}
        self.stack = []
        self.memory = {}

    def WriteToOperand(self, operand, value):
        if operand["type"] == "REG":
            self.regs[operand["reg"]] = value

        elif operand["type"] == "IMM":
            self.memory[operand["address"]] = value

        elif operand["type"] == "MEM":
            self.memory[self._get_mem_operand_target(operand)] = value

        else:
            raise DecodingError("Operand not supported")

    def ReadFromOperand(self, operand):
        """Read the operand.

        We support the following forms:

        - Immediate (IMM):  JMP 0x123456
        - Absolute Memory Address (MEM): JMP [0x123456]
        - Register (REG): JMP [EAX]
        """
        # Read from register.
        if operand["type"] == 'REG':
            return self.regs.get(operand["reg"], 0)

        # Immediate operand.
        elif operand["type"] == 'IMM':
            return operand["address"]

        # Read the content of memory.
        elif operand["type"] == "MEM":
            return self._GetMemoryAddress(
                self._get_mem_operand_target(operand), operand["size"])

        else:
            raise DecodingError("Operand not supported")

    def _get_mem_operand_target(self, operand):
        reg_base = operand["base"]
        if reg_base == "RIP":
            return operand["address"]
        else:
            # Register reference [base_reg + disp + index_reg * scale]
            return (self.regs.get(reg_base, 0) +
                    operand["disp"] +
                    self.regs.get(operand["index"], 0) * operand["scale"])

    def _GetMemoryAddress(self, offset, size):
        try:
            # First check our local cache for a previously written value.
            return self.memory[offset]
        except KeyError:
            data = self.address_space.read(offset, size)
            format_string = {1: "b", 2: "H", 4: "I", 8: "Q"}[size]

            return struct.unpack(format_string, data)[0]

    def process_lea(self, instruction):
        """Copies the address from the second operand to the first."""
        operand = instruction.operands[1]
        if operand["type"] == 'MEM':
            self.WriteToOperand(instruction.operands[0],
                                self._get_mem_operand_target(operand))
        else:
            raise DecodingError("Invalid LEA source.")

    def process_push(self, instruction):
        value = self.ReadFromOperand(instruction.operands[0])

        self.stack.append(value)

    def process_pop(self, instruction):
        try:
            value = self.stack.pop(-1)
        except IndexError:
            value = 0

        self.WriteToOperand(instruction.operands[0], value)

    def process_ret(self, _):
        if self.stack:
            return self.stack.pop(-1)

    def process_mov(self, instruction):
        value = self.ReadFromOperand(instruction.operands[1])

        self.WriteToOperand(instruction.operands[0], value)

    def process_inc(self, instruction):
        value = self.ReadFromOperand(instruction.operands[0])

        self.WriteToOperand(instruction.operands[0], value + 1)

    def process_dec(self, instruction):
        value = self.ReadFromOperand(instruction.operands[0])

        self.WriteToOperand(instruction.operands[0], value - 1)

    def process_cmp(self, instruction):
        """We dont do anything with the comparison since we dont test for it."""
        _ = instruction

    def process_test(self, instruction):
        """We dont do anything with the comparison since we dont test for it."""
        _ = instruction

    def _Operate(self, instruction, operator):
        value1 = self.ReadFromOperand(instruction.operands[0])
        value2 = self.ReadFromOperand(instruction.operands[1])

        self.WriteToOperand(
            instruction.operands[0], operator(value1, value2))

    def process_xor(self, instruction):
        return self._Operate(instruction, lambda x, y: x ^ y)

    def process_add(self, instruction):
        return self._Operate(instruction, lambda x, y: x + y)

    def process_sub(self, instruction):
        return self._Operate(instruction, lambda x, y: x - y)

    def process_and(self, instruction):
        return self._Operate(instruction, lambda x, y: x & y)

    def process_or(self, instruction):
        return self._Operate(instruction, lambda x, y: x | y)

    def process_shl(self, instruction):
        return self._Operate(instruction, lambda x, y: x << (y % 0xFF))

    def process_shr(self, instruction):
        return self._Operate(instruction, lambda x, y: x >> (y % 0xFF))

    def Inspect(self, function, instructions=10):
        """The main entry point to the Hook processor.

        We emulate the function instructions and try to determine the jump
        destination.

        Args:
           function: A basic.Function() instance.
        """
        self.Reset()
        self.address_space = function.obj_vm

        for instruction in function.disassemble(instructions=instructions):
            if instruction.is_return():
                # RET Instruction terminates processing.
                return self.process_ret(instruction)

            elif instruction.mnemonic == "call":
                return self.ReadFromOperand(instruction.operands[0])

            # A JMP instruction.
            elif instruction.is_branch():
                return self.ReadFromOperand(instruction.operands[0])

            else:
                try:
                    handler = getattr(self, "process_%s" % instruction.mnemonic)
                except AttributeError:
                    continue

                # Handle the instruction.
                try:
                    handler(instruction)
                except Exception:
                    self.session.logging.error(
                        "Unable to handle instruction %s", instruction.op_str)
                    return


class CheckPEHooks(common.WindowsCommandPlugin):
    """Checks a pe file mapped into memory for hooks."""

    name = "check_pehooks"

    __args = [
        dict(name="image_base", default=0,
             positional=True, type="SymbolAddress",
             help="The base address of the pe image in memory."),

        dict(name="type", default="all",
             choices=["all", "iat", "inline", "eat"],
             type="Choice", help="Type of hook to display."),

        dict(name="thorough", default=False, type="Boolean",
             help="By default we take some optimization. This flags forces "
             "thorough but slower checks."),
    ]

    table_header = [
        dict(name="Type", width=10),
        dict(name="source", width=20),
        dict(name="target", width=20),
        dict(name="source_func", width=60),
        dict(name="target_func"),
    ]

    def reported_access(self, address):
        """Determines if the address should be reported.

        This assesses the destination address for suspiciousness. For example if
        the address resides in a VAD region which is not mapped by a dll then it
        might be suspicious.
        """
        destination_names = self.session.address_resolver.format_address(
            address)

        # For now very simple: If any of the destination_names start with vad_*
        # it means that the address resolver cant determine which module they
        # came from.
        destination = hex(address)
        for destination in destination_names:
            if not destination.startswith("vad_"):
                return False

        return destination

    def detect_IAT_hooks(self):
        """Detect Import Address Table hooks.

        An IAT hook is where malware changes the IAT entry for a dll after its
        loaded so that when it is called from within the DLL, flow control is
        directed to the malware instead.

        We determine the IAT entry is hooked if the address is outside the dll
        which is imported.
        """
        pe = pe_vtypes.PE(image_base=self.plugin_args.image_base,
                          session=self.session)

        # First try to find all the names of the imported functions.
        imports = [
            (dll, func_name) for dll, func_name, _ in pe.ImportDirectory()]

        resolver = self.session.address_resolver

        for idx, (dll, func_address, _) in enumerate(pe.IAT()):
            try:
                target_dll, target_func_name = imports[idx]
                target_dll = self.session.address_resolver.NormalizeModuleName(
                    target_dll)
            except IndexError:
                # We can not retrieve these function's name from the
                # OriginalFirstThunk array - possibly because it is not mapped
                # in.
                target_dll = dll
                target_func_name = ""

            self.session.report_progress(
                "Checking function %s!%s", target_dll, target_func_name)

            # We only want the containing module.
            module = resolver.GetContainingModule(func_address)
            if module and target_dll == module.name:
                continue

            # Use ordinal if function has no name
            if not len(target_func_name):
                target_func_name = "(%s)" % idx

            function_name = "%s!%s" % (target_dll, target_func_name)

            # Function_name is the name which the PE file want
            yield function_name, func_address

    def collect_iat_hooks(self):
        for function_name, func_address in self.detect_IAT_hooks():
            yield dict(Type="IAT",
                       source=function_name,
                       target=utils.FormattedAddress(
                           self.session.address_resolver,
                           func_address, max_count=1, hex_if_unknown=True),
                       target_func=self.session.profile.Function(
                           func_address))

    def detect_EAT_hooks(self, size=0):
        """Detect Export Address Table hooks.

        An EAT hook is where malware changes the EAT entry for a dll after its
        loaded so that a new DLL wants to link against it, the new DLL will use
        the malware's function instead of the exporting DLL's function.

        We determine the EAT entry is hooked if the address lies outside the
        exporting dll.
        """
        address_space = self.session.GetParameter("default_address_space")
        pe = pe_vtypes.PE(image_base=self.plugin_args.image_base,
                          session=self.session,
                          address_space=address_space)
        start = self.plugin_args.image_base
        end = self.plugin_args.image_base + size

        # If the dll size is not provided we parse it from the PE header.
        if not size:
            for _, _, virtual_address, section_size in pe.Sections():
                # Only count executable sections.
                section_end = (self.plugin_args.image_base +
                               virtual_address + section_size)
                if section_end > end:
                    end = section_end

        resolver = self.session.address_resolver

        for dll, func, name, hint in pe.ExportDirectory():
            self.session.report_progress("Checking export %s!%s", dll, name)

            # Skip zero or invalid addresses.
            if address_space.read(func.v(), 10) == b"\x00" * 10:
                continue

            # Report on exports which fall outside the dll.
            if start < func.v() < end:
                continue

            function_name = "%s:%s (%s)" % (
                resolver.NormalizeModuleName(dll), name, hint)

            yield function_name, func

    def collect_eat_hooks(self):
        for function_name, func_address in self.detect_EAT_hooks():
            yield dict(Type="EAT",
                       source=function_name,
                       target=utils.FormattedAddress(
                           self.session.address_resolver,
                           func_address, max_count=1, hex_if_unknown=True),
                       target_func=self.session.profile.Function(
                           func_address))

    def detect_inline_hooks(self):
        """A Generator of hooked exported functions from this PE file.

        Yields:
          A tuple of (function, name, jump_destination)
        """
        # Inspect the export directory for inline hooks.
        pe = pe_vtypes.PE(image_base=self.plugin_args.image_base,
                          address_space=self.session.GetParameter(
                              "default_address_space"),
                          session=self.session)
        pfn_db = self.session.profile.get_constant_object("MmPfnDatabase")
        heuristic = HookHeuristic(session=self.session)

        ok_pages = set()

        for _, function, name, _ in pe.ExportDirectory():
            # Dereference the function pointer.
            function_address = function.deref().obj_offset

            self.session.report_progress(
                "Checking function %#x (%s)", function, name)

            # Check if the page is private or a file mapping. Usually if a
            # mapped page is modified it will be converted to a private page due
            # to Windows copy on write semantics. We assume that hooks are only
            # placed in memory, and therefore functions which are still mapped
            # to disk files are not hooked and can be safely skipped.
            if not self.plugin_args.thorough:
                # We must do the vtop in the process address space. This is the
                # physical page backing the function preamble.
                phys_address = function.obj_vm.vtop(function_address)

                # Page not mapped.
                if phys_address == None:
                    continue

                phys_page = phys_address >> 12

                # We determined this page is ok before - we can skip it.
                if phys_page in ok_pages:
                    continue

                # Get the PFN DB record.
                pfn_obj = pfn_db[phys_page]

                # The page is controlled by a prototype PTE which means it is
                # still a file mapping. It has not been changed.
                if pfn_obj.IsPrototype:
                    ok_pages.add(phys_page)
                    continue

            # Try to detect an inline hook.
            destination = heuristic.Inspect(function, instructions=3) or ""

            # If we did not detect a hook we skip this function.
            if destination:
                yield function, name, destination

    def collect_inline_hooks(self):
        for function, _, destination in self.detect_inline_hooks():
            hook_detected = False

            # Try to resolve the destination into a name.
            destination_name = self.reported_access(destination)

            # We know about it. We suppress the output for jumps that go into a
            # known module. These should be visible using the regular vad
            # module.
            if destination_name:
                hook_detected = True

            # Skip non hooked results if verbosity is too low.
            if self.plugin_args.verbosity < 10 and not hook_detected:
                continue

            # Only highlight results if verbosity is high.
            highlight = ""
            if hook_detected and self.plugin_args.verbosity > 1:
                highlight = "important"

            yield dict(Type="Inline",
                       source=utils.FormattedAddress(
                           self.session.address_resolver,
                           function.deref(), max_count=1),
                       target=utils.FormattedAddress(
                           self.session.address_resolver,
                           destination, max_count=1),
                       source_func=function.deref(),
                       target_func=self.session.profile.Function(
                           destination),
                       highlight=highlight)

    def collect(self):
        if self.plugin_args.type in ["all", "inline"]:
            for x in self.collect_inline_hooks():
                yield x

        if self.plugin_args.type in ["all", "iat"]:
            for x in self.collect_iat_hooks():
                yield x

        if self.plugin_args.type in ["all", "eat"]:
            for x in self.collect_eat_hooks():
                yield x


class EATHooks(common.WinProcessFilter):
    """Detect EAT hooks in process and kernel memory"""

    name = "hooks_eat"

    table_header = [
        dict(name="divider", type="Divider"),
        dict(name="_EPROCESS", hidden=True),
        dict(name="Type", hidden=True),
        dict(name="source", width=20),
        dict(name="target", width=20),
        dict(name="target_func"),
    ]

    checker_method = CheckPEHooks.collect_eat_hooks

    def column_types(self):
        return dict(_EPROCESS=self.session.profile._EPROCESS,
                    source="",
                    target="",
                    target_func=self.session.profile.Function())

    def collect_hooks(self, task, dll):
        checker = self.session.plugins.check_pehooks(
            image_base=dll.base, thorough=self.plugin_args.thorough)

        for info in self.__class__.checker_method(checker):
            info["_EPROCESS"] = task
            yield info

    def collect(self):
        cc = self.session.plugins.cc()
        with cc:
            for task in self.filter_processes():
                cc.SwitchProcessContext(task)

                yield dict(divider="Process %s (%s)" % (task.name, task.pid))

                for dll in task.get_load_modules():
                    for x in self.collect_hooks(task, dll):
                        yield x


class TestEATHooks(testlib.SimpleTestCase):
    PLUGIN = "hooks_eat"

    PARAMETERS = dict(
        commandline="hooks_eat %(pids)s"
        )


class IATHooks(EATHooks):
    """Detect IAT/EAT hooks in process and kernel memory"""

    name = "hooks_iat"
    checker_method = CheckPEHooks.collect_iat_hooks


class TestIATHooks(testlib.SimpleTestCase):
    PLUGIN = "hooks_iat"

    PARAMETERS = dict(
        commandline="hooks_iat  %(pids)s"
        )


class InlineHooks(EATHooks):
    """Detect API hooks in process and kernel memory"""

    name = "hooks_inline"
    checker_method = CheckPEHooks.collect_inline_hooks
    table_header = [
        dict(name="divider", type="Divider"),
        dict(name="_EPROCESS", hidden=True),
        dict(name="source", width=20),
        dict(name="target", width=20),
        dict(name="Type", hidden=True),
        dict(name="source_func", width=60),
        dict(name="target_func"),
    ]


class TestInlineHooks(testlib.SimpleTestCase):
    PLUGIN = "hooks_inline"

    PARAMETERS = dict(
        commandline="hooks_inline %(pids)s"
        )
