# Volatility
# Copyright (c) 2010, 2011, 2012 Michael Ligh <michael.ligh@mnin.org>
#
# 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 volatility.utils as utils
import volatility.obj as obj
import volatility.plugins.common as common
import volatility.debug as debug
import volatility.win32.tasks as tasks
import volatility.win32.modules as modules

#--------------------------------------------------------------------------------
# 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']],
    }]}

#--------------------------------------------------------------------------------
# profile modifications 
#--------------------------------------------------------------------------------

class MalwareTimerVTypes(obj.ProfileModification):
    before = ['WindowsOverlay']
    conditions = {'os': lambda x: x == 'windows',
                  'memory_model': lambda x: x == '32bit'}
    def modification(self, profile):
        profile.vtypes.update(timer_types)

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

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

    @staticmethod
    def is_valid_profile(profile):
        return (profile.metadata.get('os', 'unknown') == 'windows' and
                profile.metadata.get('memory_model', '32bit') == '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.zread(func_addr, 200)

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

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

    def calculate(self):
        addr_space = utils.load_as(self._config)

        if not self.is_valid_profile(addr_space.profile):
            debug.error("This command does not support the selected profile.")

        # Get the OS version we're analyzing 
        version = (addr_space.profile.metadata.get('major', 0),
                   addr_space.profile.metadata.get('minor', 0))

        modlist = list(modules.lsmod(addr_space))
        mods = dict((mod.DllBase, mod) for mod in modlist)
        mod_addrs = sorted(mods.keys())

        # 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
                                addr_space.profile.metadata.get('build', 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 = obj.Object("Array", offset = KiTimerTableListHead,
                                        vm = addr_space,
                                        targetType = '_LIST_ENTRY',
                                        count = 256)

            for l in lists:
                for t in l.list_of_type("_KTIMER", "TimerListEntry"):
                    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 = obj.Object("Array", offset = KiTimerTableListHead,
                                        vm = addr_space,
                                        targetType = '_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).
            for kpcr in tasks.get_kdbg(addr_space).kpcrs():
                for table in kpcr.ProcessorBlock.TimerTable.TimerEntries:
                    for t in table.Entry.list_of_type("_KTIMER", "TimerListEntry"):
                        timers.append(t)

        for timer in 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 = tasks.find_module(mods, mod_addrs, timer.Dpc.DeferredRoutine)

            yield timer, module

    def render_text(self, outfd, data):

        self.table_header(outfd,
                        [("Offset(V)", "[addrpad]"),
                         ("DueTime", "24"),
                         ("Period(ms)", "10"),
                         ("Signaled", "10"),
                         ("Routine", "[addrpad]"),
                         ("Module", ""),
                        ])

        for timer, module in data:

            if timer.Header.SignalState.v():
                signaled = "Yes"
            else:
                signaled = "-"

            if module:
                module_name = str(module.BaseDllName or '')
            else:
                module_name = "UNKNOWN"

            due_time = "{0:#010x}:{1:#010x}".format(timer.DueTime.HighPart, timer.DueTime.LowPart)

            self.table_row(outfd,
                        timer.obj_offset,
                        due_time,
                        timer.Period,
                        signaled,
                        timer.Dpc.DeferredRoutine,
                        module_name)

