# 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
#
import re

from rekall import plugin
from rekall import utils
from rekall.plugins.windows import filescan

class DeviceTree(filescan.DriverScan):
    "Show device tree."

    __name = "devicetree"

    def render(self, renderer):
        renderer.table_header([
            dict(name="Type", type="TreeNode", width=10, max_depth=10),
            dict(name="Address", style="address", padding="0"),
            dict(name="Name", width=20),
            dict(name="Device Type", width=30),
            dict(name="Path"),
        ])

        for _, _, driver_obj, _, _ in self.generate_hits():
            renderer.table_row(
                utils.AttributedString("DRV", [(0, 30, "BLACK", "RED")]),
                driver_obj.obj_offset,
                driver_obj.DriverName.v(vm=self.kernel_address_space),
                depth=0)

            first_device = driver_obj.DeviceObject.dereference(
                vm=self.kernel_address_space)

            for device in first_device.walk_list("NextDevice"):
                device_header = self.profile.Object(
                    "_OBJECT_HEADER", offset=device.obj_offset -
                    device.obj_profile.get_obj_offset("_OBJECT_HEADER", "Body"),
                    vm=device.obj_vm)

                device_name = device_header.NameInfo.Name.cast(
                    vm=self.kernel_address_space)

                renderer.table_row(
                    utils.AttributedString("DEV", [(0, 30, "WHITE", "BLUE")]),
                    device.obj_offset, device_name,
                    device.DeviceType, depth=1)

                level = 1

                for att_device in device.walk_list("AttachedDevice"):
                    renderer.table_row(
                        utils.AttributedString("ATT", [(0, 30, "BLACK", "GREEN")]),
                        att_device.obj_offset,
                        device_name, att_device.DeviceType,
                        att_device.DriverObject.DriverName,
                        depth=level)

                    level += 1


class DriverIrp(plugin.VerbosityMixIn, filescan.DriverScan):
    "Driver IRP hook detection"

    __name = "driverirp"

    mod_re = None

    @classmethod
    def args(cls, parser):
        """Declare the command line args we need."""
        super(DriverIrp, cls).args(parser)
        parser.add_argument("-r", "--regex",
                            help='Analyze drivers matching REGEX')

    def __init__(self, regex=None, **kwargs):
        super(DriverIrp, self).__init__(**kwargs)
        if regex:
            self.mod_re = re.compile(regex, re.I)

    def render(self, renderer):
        module_plugin = self.session.plugins.modules()

        for _, _, driver_obj, _, driver_name in self.generate_hits():
            # Continue if a regex was supplied and it doesn't match
            if self.mod_re:
                if not driver_name:
                    continue

                # Continue if a regex was supplied and it doesn't match
                if not self.mod_re.search(driver_name):
                    continue

            # Write the standard header for each driver object
            renderer.section()

            renderer.format("DriverName: {0}\n", driver_name)
            renderer.format("DriverStart: {0:addrpad}\n",
                            driver_obj.DriverStart)

            renderer.format("DriverSize: {0:addrpad}\n", driver_obj.DriverSize)
            renderer.format("DriverStartIo: {0:addrpad}\n",
                            driver_obj.DriverStartIo)

            renderer.table_header([('', "index", ">4"),
                                   ('Func Name', "function", "36"),
                                   ('Func Addr', "func_addres", '[addrpad]'),
                                   ('Module', "name", '')
                                  ])

            # Write the address and owner of each IRP function
            for i, function in enumerate(driver_obj.MajorFunction):
                # Make sure this is in the kernel address space.
                function = driver_obj.MajorFunction[i].dereference(
                    vm=self.kernel_address_space)

                # Which module does this function come from?
                module = module_plugin.find_module(function)

                renderer.table_row(i, function.obj_name, function.obj_offset,
                                   module.FullDllName)

                # For verbose we print a bit about the function
                # (e.g. disassembly if available.).
                if self.verbosity > 1:
                    renderer.format("{0}\n\n", function)
