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

#--------------------------------------------------------------------------------
# vtypes
#--------------------------------------------------------------------------------

# This type is defined in Win2K3SP0x86 and VistaSP2x86, but
# it applies to many other profiles in which it is not defined
# in the public PDBs.
timer_types = {
    '_KTIMER_TABLE_ENTRY' : [ 0x10, {
    'Entry' : [ 0x0, ['_LIST_ENTRY']],
    'Time'  : [ 0x8, ['_ULARGE_INTEGER']],
    }]}

#--------------------------------------------------------------------------------
# timers
#--------------------------------------------------------------------------------

class Timers(common.WindowsCommandPlugin):
    """Print kernel timers and associated module DPCs"""

    __name = "timers"

    def __init__(self, **kwargs):
        """Print kernel timers and associated module DPCs."""
        super(Timers, self).__init__(**kwargs)
        self.profile.add_types(timer_types)

    @classmethod
    def is_active(cls, config):
        """Timers are currently only supported on 32 bit windows."""
        profile = config.profile

        return (profile and profile.metadata("os") == 'windows' and
                profile.metadata("memory_model") == "32bit")

    def find_list_head(self, nt_mod, func, sig):
        """
        Find the KiTimerTableListHead given an exported
        function as a starting point and a small signature.

        @param nt_mod: _LDR_DATA_TABLE_ENTRY object for NT module
        @param func: function name exported by the NT module
        @param sig: byte string/pattern to use for finding the symbol
        """

        # Lookup the exported function
        func_rva = nt_mod.getprocaddress(func)
        if func_rva == None:
            return None

        func_addr = func_rva + nt_mod.DllBase

        # Read enough of the function prolog
        data = nt_mod.obj_vm.read(func_addr, 200)

        # Scan for the byte signature
        n = data.find(sig)
        if n == -1:
            return None

        return self.profile.Object('address', func_addr + n + len(sig), nt_mod.obj_vm)

    def calculate(self):
        # Get the OS version we're analyzing
        version = (self.profile.metadata('major', 0),
                   self.profile.metadata('minor', 0))

        modules = self.session.plugins.modules(session=self.session)
        mod_addrs = modules.addresses()
        modlist = list(modules.lsmod())

        # KTIMERs collected
        timers = []

        # Valid KTIMER.Header.Type values
        TimerNotificationObject    = 8
        TimerSynchronizationObject = 9
        valid_types = (TimerNotificationObject, TimerSynchronizationObject)

        if version == (5,1) or (version == (5,2) and
                                self.profile.metadata.get('major', 0) == 3789):
            """
            On XP SP0-SP3 x86 and Windows 2003 SP0, KiTimerTableListHead
            is an array of 256 _LIST_ENTRY for _KTIMERs.
            """
            KiTimerTableListHead = self.find_list_head(modlist[0],
                                        "KeUpdateSystemTime",
                                        "\x25\xFF\x00\x00\x00\x8D\x0C\xC5")

            lists = self.profile.Object("Array", offset = KiTimerTableListHead,
                                        vm = self.kernel_address_space,
                                        target = '_LIST_ENTRY',
                                        count = 256)

            for l in lists:
                for t in l.list_of_type("_KTIMER", "TimerListEntry"):
                    if t not in timers:
                        timers.append(t)

        elif version == (5,2) or version == (6,0):
            """
            On XP x64, Windows 2003 SP1-SP2, and Vista SP0-SP2, KiTimerTableListHead
            is an array of 512 _KTIMER_TABLE_ENTRY structs.
            """
            KiTimerTableListHead = self.find_list_head(modlist[0],
                                        "KeCancelTimer",
                                        "\xC1\xE7\x04\x81\xC7")

            lists = self.profile.Object("Array", offset = KiTimerTableListHead,
                                        vm = self.kernel_address_space,
                                        target = '_KTIMER_TABLE_ENTRY',
                                        count = 512)

            for l in lists:
                for t in l.Entry.list_of_type("_KTIMER", "TimerListEntry"):
                    timers.append(t)

        elif version == (6,1):
            """
            On Windows 7, there is no more KiTimerTableListHead. The list is
            at _KPCR.PrcbData.TimerTable.TimerEntries (credits to Matt Suiche
            for this one. See http://pastebin.com/FiRsGW3f).
            """
            kpcrscan = self.session.plugins.kpcrscan(session=self.session)
            for kpcr in kpcrscan.generate_hits():
                # Note that if this is not the correct kpcr, the NoneObject will
                # break out of the following loop:
                for table in kpcr.PrcbData.TimerTable.TimerEntries:
                    for t in table.Entry.list_of_type("_KTIMER", "TimerListEntry"):
                        timers.append(t)
        else:
            raise RuntimeError("Timers are not supported on this Profile.")

        # Sort the timers by address to make them easier to inspect.
        for timer in sorted(timers):
            # Sanity check on the timer type
            if timer.Header.Type not in valid_types:
                continue

            # Ignore timers without DPCs
            if not timer.Dpc.is_valid() or not timer.Dpc.DeferredRoutine.is_valid():
                continue

            # Lookup the module containing the DPC
            module = modules.find_module(timer.Dpc.DeferredRoutine)

            yield timer, module

    def render(self, renderer):
        renderer.table_header([("Offset", "offset", "[addrpad]"),
                               ("DueTime(H)          ", "due_high", "#020x"),
                               ("Period(ms)", "period", ">10"),
                               ("Signaled", "signaled", ">10"),
                               ("Routine", "routine", "[addrpad]"),
                               ("Module", "module", "20")])

        for timer, module in self.calculate():
            if timer.Header.SignalState.v():
                signaled = "Yes"
            else:
                signaled = "-"

            if module:
                module_name = module.BaseDllName
            else:
                module_name = "UNKNOWN"

            renderer.table_row(
                timer,
                timer.DueTime.QuadPart,
                timer.Period,
                signaled,
                timer.Dpc.DeferredRoutine,
                module_name)
