# 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 obj
from rekall.plugins.windows import filescan

class _DEVICE_OBJECT(obj.Struct):
    "Class for device objects"

    def next_devices(self):
        """Iterate over all the devices in the NextDevice entry."""
        while self:
            yield self
            self = self.NextDevice.dereference()

    def attached_devices(self):
        "Enumerate the device's attachees"
        device = self.AttachedDevice.dereference()
        while device:
            yield device
            device = device.AttachedDevice.dereference()


class DeviceModification(obj.ProfileModification):
    @classmethod
    def modify(cls, profile):
        profile.add_classes({
                '_DEVICE_OBJECT': _DEVICE_OBJECT,
                })


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

    __name = "devicetree"

    def __init__(self, **kwargs):
        super(DeviceTree, self).__init__(**kwargs)
        self.profile = DeviceModification(self.profile)

    def render(self, renderer):
        for _, _, driver_obj, _, _ in self.generate_hits():
            renderer.format(
                "DRV 0x{0:08x} {1}\n", driver_obj.obj_offset,
                driver_obj.DriverName.v(vm=self.kernel_address_space))

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

            for device in first_device.next_devices():
                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.v(
                    vm=self.kernel_address_space)

                renderer.format("---| DEV {0:#x} {1} {2}\n", device.obj_offset,
                                device_name, device.DeviceType)

                level = 0

                for att_device in device.attached_devices():
                    renderer.format(
                        "------{0}| ATT {1:#x} {2} - {3} {4}\n",
                        "---" * level,
                        att_device.obj_offset,
                        device_name,
                        att_device.DriverObject.DriverName,
                        att_device.DeviceType)

                    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)
        self.profile = DeviceModification(self.profile)
        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:#x}\n", driver_obj.DriverStart)
            renderer.format("DriverSize: {0:#x}\n", driver_obj.DriverSize)
            renderer.format("DriverStartIo: {0:#x}\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)
