# Rekall Memory Forensics
# Copyright (c) 2010, 2011, 2012 Michael Ligh <michael.ligh@mnin.org>
# Copyright 2013 Google Inc. All Rights Reserved.
#
# 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 str
from rekall import obj
from rekall import testlib
from rekall.plugins import core
from rekall.plugins.windows import common

# pylint: disable=protected-access


class Malfind(core.DirectoryDumperMixin, common.WinProcessFilter):
    "Find hidden and injected code"

    __name = "malfind"

    dump_dir_optional = True
    default_dump_dir = None

    def _is_vad_empty(self, vad, address_space):
        """
        Check if a VAD region is either entirely unavailable
        due to paging, entirely consiting of zeros, or a
        combination of the two. This helps ignore false positives
        whose VAD flags match task._injection_filter requirements
        but there's no data and thus not worth reporting it.

        @param vad: an MMVAD object in kernel AS
        @param address_space: the process address space
        """

        PAGE_SIZE = 0x1000
        all_zero_page = "\x00" * PAGE_SIZE

        offset = 0
        while offset < vad.Length:
            next_addr = vad.Start + offset
            if (address_space.is_valid_address(next_addr) and
                    address_space.read(next_addr, PAGE_SIZE) != all_zero_page):
                return False
            offset += PAGE_SIZE

        return True

    def _injection_filter(self, vad, task_as):
        """Detects injected vad regions.

        This looks for private allocations that are committed,
        memory-resident, non-empty (not all zeros) and with an
        original protection that includes write and execute.

        It is important to note that protections are applied at
        the allocation granularity (page level). Thus the original
        protection might not be the current protection, and it
        also might not apply to all pages in the VAD range.

        @param vad: an MMVAD object.

        @returns: True if the MMVAD looks like it might
        contain injected code.
        """
        # Try to find injections.
        protect = str(vad.u.VadFlags.ProtectionEnum)
        write_exec = "EXECUTE" in protect and "WRITE" in protect

        # The Write/Execute check applies to everything
        if not write_exec:
            return False

        # This is a typical VirtualAlloc'd injection
        if ((vad.u.VadFlags.PrivateMemory == 1 and
             vad.Tag == "VadS") or

                # This is a stuxnet-style injection
                (vad.u.VadFlags.PrivateMemory == 0 and
                 protect != "EXECUTE_WRITECOPY")):
            return not self._is_vad_empty(vad, task_as)

        return False

    def render(self, renderer):
        cc = self.session.plugins.cc()
        for task in self.filter_processes():
            task_as = task.get_process_address_space()
            if not task_as:
                continue

            with cc:
                cc.SwitchProcessContext(task)
                for vad in task.RealVadRoot.traverse():
                    self.session.report_progress("Checking %r of pid %s",
                                                 vad, task.UniqueProcessId)

                    if self._injection_filter(vad, task_as):
                        renderer.section()
                        renderer.format(
                            "Process: {0} Pid: {1} Address: {2:#x}\n",
                            task.ImageFileName, task.UniqueProcessId, vad.Start)

                        renderer.format("Vad Tag: {0} Protection: {1}\n",
                                        vad.Tag, vad.u.VadFlags.ProtectionEnum)

                        renderer.format("Flags: {0}\n", vad.u.VadFlags)
                        renderer.format("\n")

                        dumper = self.session.plugins.dump(
                            offset=vad.Start, rows=4)
                        dumper.render(renderer, suppress_headers=True)

                        renderer.format("\n")

                        disassembler = self.session.plugins.dis(
                            offset=vad.Start, length=0x40)
                        disassembler.render(renderer, suppress_headers=True)

                        if self.dump_dir:
                            filename = "{0}.{1:d}.{2:08x}-{3:08x}.dmp".format(
                                task.ImageFileName, task.pid, vad.Start,
                                vad.End)

                            with renderer.open(directory=self.dump_dir,
                                               filename=filename,
                                               mode='wb') as fd:
                                self.session.report_progress(
                                    "Dumping %s" % filename)

                                self.CopyToFile(task_as, vad.Start, vad.End, fd)


class LdrModules(common.WinProcessFilter):
    "Detect unlinked DLLs"

    __name = "ldrmodules"

    table_header = [
        dict(name="divider", type="Divider"),
        dict(name="_EPROCESS", hidden=True),
        dict(name="base", style="address"),
        dict(name="in_load", width=5),
        dict(name="in_load_path", width=80, hidden=True),
        dict(name="in_init", width=5),
        dict(name="in_init_path", width=80, hidden=True),
        dict(name="in_mem", width=5),
        dict(name="in_mem_path", width=80, hidden=True),
        dict(name="mapped")
    ]

    def column_types(self):
        return dict(
            _EPROCESS=self.session.profile._EPROCESS(),
            base_address=0,
            in_load=False,
            in_load_path=self.session.profile._UNICODE_STRING(),
            in_init=False,
            in_init_path=self.session.profile._UNICODE_STRING(),
            in_mem=False,
            in_mem_path=self.session.profile._UNICODE_STRING(),
            mapped_filename=self.session.profile._UNICODE_STRING(),
        )

    def list_mapped_files(self, task):
        """Iterates over all vads and returns executable regions.

        Yields:
          vad objects which are both executable and have a file name.
        """
        self.session.report_progress("Inspecting Pid %s",
                                     task.UniqueProcessId)

        for vad in task.RealVadRoot.traverse():
            try:
                file_obj = vad.ControlArea.FilePointer
                protect = str(vad.u.VadFlags.ProtectionEnum)
                if "EXECUTE" in protect and "WRITE" in protect:
                    yield vad, file_obj.FileName
            except AttributeError:
                pass


    def collect(self):
        for task in self.filter_processes():
            # Build a dictionary for all three PEB lists where the
            # keys are base address and module objects are the values
            inloadorder = dict((mod.DllBase.v(), mod)
                               for mod in task.get_load_modules())

            ininitorder = dict((mod.DllBase.v(), mod)
                               for mod in task.get_init_modules())

            inmemorder = dict((mod.DllBase.v(), mod)
                              for mod in task.get_mem_modules())

            # Build a similar dictionary for the mapped files
            mapped_files = dict((vad.Start, name)
                                for vad, name in self.list_mapped_files(task))

            yield dict(divider=task)

            # For each base address with a mapped file, print info on
            # the other PEB lists to spot discrepancies.
            for base in list(mapped_files.keys()):
                yield dict(_EPROCESS=task,
                           base=base,
                           in_load=base in inloadorder,
                           in_load_path=inloadorder.get(
                               base, obj.NoneObject()).FullDllName,
                           in_init=base in ininitorder,
                           in_init_path=ininitorder.get(
                               base, obj.NoneObject()).FullDllName,
                           in_mem=base in inmemorder,
                           in_mem_path=inmemorder.get(
                               base, obj.NoneObject()).FullDllName,
                           mapped=mapped_files[base])


class TestLdrModules(testlib.SimpleTestCase):
    PARAMETERS = dict(
        commandline="ldrmodules %(pids)s"
        )
