# Volatility
# Copyright (C) 2007-2013 Volatility Foundation
# Copyright (c) 2010, 2011, 2012 Michael Ligh <michael.ligh@mnin.org>
#
# This file is part of Volatility.
#
# Volatility is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License Version 2 as
# published by the Free Software Foundation.  You may not use, modify or
# distribute this program under any other version of the GNU General
# Public License.
#
# Volatility 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 Volatility.  If not, see <http://www.gnu.org/licenses/>.
#

import sys, pydoc
import volatility.utils as utils
import volatility.registry as registry
import volatility.obj as obj
import volatility.win32.modules as modules
import volatility.win32.tasks as tasks
import volatility.plugins.ssdt as ssdt
import volatility.plugins.taskmods as taskmods
import volatility.plugins.modscan as modscan
import volatility.plugins.malware.malfind as malfind
import volatility.debug as debug

try:
    import distorm3 #pylint: disable-msg=W0611
    has_distorm3 = True
except ImportError:
    has_distorm3 = False

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

thread_types = {
    '_KTHREAD' : [ None , {
    'State' : [ None, ['Enumeration', dict(target = 'unsigned char', choices = {
        0: 'Initialized', 1: 'Ready', 2: 'Running', 3: 'Standby', 4: 'Terminated',
        5: 'Waiting', 6: 'Transition', 7: 'DeferredReady', 8: 'GateWait'})]],
    'WaitReason' : [ None, ['Enumeration', dict(target = 'unsigned char', choices = {
        0: 'Executive', 1: 'FreePage', 2: 'PageIn', 3: 'PoolAllocation',
        4: 'DelayExecution', 5: 'Suspended', 6: 'UserRequest', 7: 'WrExecutive',
        8: 'WrFreePage', 9: 'WrPageIn', 10: 'WrPoolAllocation', 11: 'WrDelayExecution',
        12: 'WrSuspended', 13: 'WrUserRequest', 14: 'WrEventPair', 15: 'WrQueue',
        16: 'WrLpcReceive', 17: 'WrLpcReply', 18: 'WrVirtualMemory', 19: 'WrPageOut',
        20: 'WrRendezvous', 21: 'Spare2', 22: 'Spare3', 23: 'Spare4', 24: 'Spare5',
        25: 'Spare6', 26: 'WrKernel', 27: 'WrResource', 28: 'WrPushLock', 29: 'WrMutex',
        30: 'WrQuantumEnd', 31: 'WrDispatchInt', 32: 'WrPreempted',
        33: 'WrYieldExecution', 34: 'WrFastMutex', 35: 'WrGuardedMutex',
        36: 'WrRundown', 37: 'MaximumWaitReason'})]],
    }],
    '_ETHREAD': [ None, {
    'CrossThreadFlags': [ None, ['Flags', {'bitmap': {
        'PS_CROSS_THREAD_FLAGS_TERMINATED': 0,
        'PS_CROSS_THREAD_FLAGS_DEADTHREAD': 1,
        'PS_CROSS_THREAD_FLAGS_HIDEFROMDBG': 2,
        'PS_CROSS_THREAD_FLAGS_IMPERSONATING': 3,
        'PS_CROSS_THREAD_FLAGS_SYSTEM': 4,
        'PS_CROSS_THREAD_FLAGS_HARD_ERRORS_DISABLED': 5,
        'PS_CROSS_THREAD_FLAGS_BREAK_ON_TERMINATION': 6,
        'PS_CROSS_THREAD_FLAGS_SKIP_CREATION_MSG': 7,
        'PS_CROSS_THREAD_FLAGS_SKIP_TERMINATION_MSG': 8,
        }}]],
    }],
}

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

class MalwareKthread(obj.ProfileModification):
    before = ['WindowsObjectClasses', 'WindowsOverlay']
    conditions = {'os': lambda x: x == 'windows'}
    def modification(self, profile):
        profile.merge_overlay(thread_types)

#--------------------------------------------------------------------------------
# thread checks
#--------------------------------------------------------------------------------

class AbstractThreadCheck(object):
    """Base thread check class"""

    def __init__(self, thread, mods, mod_addrs, \
                    hooked_tables, found_by_scanner):
        """
        @param thread: the _ETHREAD object

        @param mods: a dictionary with module bases as 
        keys and _LDR_DATA_TABLE_ENTRY as values. 

        @param mod_addrs: a sorted list of module base
        addresses 

        @param hooked_tables: a list of SSDTs that have
        one or more hooked functions. 

        @param found_by_scanner: True/False if the _ETHREAD
        passed as the thread parameter was found via 
        list walking or pool scanning. 
        """
        self.thread = thread
        self.mods = mods
        self.mod_addrs = mod_addrs
        self.hooked_tables = hooked_tables
        self.found_by_scanner = found_by_scanner
        self.flags = str(thread.CrossThreadFlags)

    def check(self):
        """Return True or False from this method"""

class OrphanThread(AbstractThreadCheck):
    """Detect orphan threads"""

    def check(self):
        """This check is True for system threads whose start address
        do not map back to known/loaded kernel drivers."""

        # Take the address space from any module object
        addr_space = self.mods.values()[0].obj_vm

        module = tasks.find_module(self.mods,
            self.mod_addrs, addr_space.address_mask(self.thread.StartAddress))

        return ('PS_CROSS_THREAD_FLAGS_SYSTEM' in self.flags and
                    module == None)

class DkomExit(AbstractThreadCheck):
    """Detect inconsistencies wrt exit times and termination"""

    def check(self):
        """This check is True when a thread's ExitTime is non-zero
        (indicating it has exited) but the state and flags 
        indicate that it is still active."""

        return (self.thread.ExitTime != 0 and
                    str(self.thread.Tcb.State) != 'Terminated' and
                    not 'PS_CROSS_THREAD_FLAGS_TERMINATED' in self.flags)

class HideFromDebug(AbstractThreadCheck):
    """Detect threads hidden from debuggers"""

    def check(self):
        """This check is True when a thread's flags report that
        it is being hidden from a debugger."""

        return 'PS_CROSS_THREAD_FLAGS_HIDEFROMDBG' in self.flags

class SystemThread(AbstractThreadCheck):
    """Detect system threads"""

    def check(self):
        """This check is True when a thread's flags report that 
        it is a system thread (i.e. PsCreateSystemThread)."""

        return 'PS_CROSS_THREAD_FLAGS_SYSTEM' in self.flags

class Impersonation(AbstractThreadCheck):
    """Detect impersonating threads"""

    def check(self):
        """This check is True when a thread's flags indicate that 
        it is impersonating another thread's security context."""

        return 'PS_CROSS_THREAD_FLAGS_IMPERSONATING' in self.flags

class HwBreakpoint(AbstractThreadCheck):
    """Detect threads with hardware breakpoints"""

    def check(self):
        """This check is True when a thread's trap frame shows
        usage of the Dr* registers in a manner consistent with
        hardware breakpoints."""

        # Don't check threads that appear to have exited
        if self.found_by_scanner:
            return False

        if 'PS_CROSS_THREAD_FLAGS_TERMINATED' in self.flags:
            return False

        trap = self.thread.Tcb.TrapFrame.dereference_as("_KTRAP_FRAME")

        if not trap:
            return False

        if ((trap.Dr0 != 0 or trap.Dr1 != 0 or trap.Dr2 != 0
                or trap.Dr3 != 0) and
                (trap.Dr6 != 0 and trap.Dr7 != 0)):
            return True

        return False

class AttachedProcess(AbstractThreadCheck):
    """Detect threads attached to another process"""

    def check(self):
        """This check is True when a thread is currently attached
        to a process other than the process that owns the thread."""

        return (self.thread.ExitTime == 0 and
                self.thread.owning_process().obj_offset !=
                self.thread.attached_process().obj_offset)

class HookedSSDT(AbstractThreadCheck):
    """Check if a thread is using a hooked SSDT"""

    def check(self):
        """This check is True if any of the thread's SSDTs have 
        hooked functions. If its True and the SSDT hooking module
        is legit, you can filter them out with --allow-hook."""

        # Check doesn't apply to x64
        if self.hooked_tables == None:
            return False

        ssdt_obj = self.thread.Tcb.ServiceTable.\
            dereference_as('_SERVICE_DESCRIPTOR_TABLE')

        for _, desc in enumerate(ssdt_obj.Descriptors):
            table = desc.KiServiceTable.v()
            if table in self.hooked_tables.keys():
                return True
        return False

class ScannerOnly(AbstractThreadCheck):
    """Detect threads no longer in a linked list"""

    def check(self):
        """This check is True when a thread is found by pool tag
        scanning but not in list traversal."""

        return self.found_by_scanner

#--------------------------------------------------------------------------------
# threads plugin 
#--------------------------------------------------------------------------------

class Threads(taskmods.DllList):
    "Investigate _ETHREAD and _KTHREADs"

    def __init__(self, config, *args, **kwargs):
        taskmods.DllList.__init__(self, config, *args, **kwargs)
        self.bits32 = None
        config.add_option("FILTER", short_option = 'F', default = None,
                            help = 'Tags to filter (comma-separated)')

        config.add_option("LISTTAGS", short_option = 'L', default = False,
                          action = 'store_true', help = 'List all available tags')

    def get_hooked_tables(self, addr_space):
        """This function finds SSDTs in an address space, checks
        if there are any hooked functions in the SSDTs, and returns
        a dictionary where SSDT base addresses are the keys and the
        values are lists of hooked function names.

        @param addr_space: a kernel address space. 
        """

        # Names of the legit executive modules for SSDT tables 
        executive_modules = [
            # SSDT 0
            ["ntoskrnl.exe", "ntkrnlpa.exe", "ntkrnlmp.exe", "ntkrpamp.exe"],
            # SSDT 1 
            ["win32k.sys"],
            # SSDT 2
            ["spud.sys"],
            # SSDT 3
            []]

        syscalls = addr_space.profile.syscalls

        hooked_tables = {}

        for info in ssdt.SSDT(self._config).calculate():
            idx, table, n, vm, mods, mod_addrs = info
            # This is straight out of ssdt.py. Too bad there's no better way 
            # to not duplicate code?
            for i in range(n):
                if self.bits32:
                    # These are absolute function addresses in kernel memory. 
                    syscall_addr = obj.Object('address', table + (i * 4), vm).v()
                else:
                    # These must be signed long for x64 because they are RVAs 
                    # relative to the base of the table and can be negative. 
                    offset = obj.Object('long', table + (i * 4), vm).v()
                    # The offset is the top 20 bits of the 32 bit number. 
                    syscall_addr = table + (offset >> 4)
                try:
                    syscall_name = syscalls[idx][i]
                except IndexError:
                    syscall_name = "UNKNOWN"

                syscall_mod = tasks.find_module(mods, mod_addrs, syscall_addr)
                if syscall_mod:
                    syscall_modname = syscall_mod.BaseDllName
                else:
                    syscall_modname = "UNKNOWN"

                if str(syscall_modname).lower() not in executive_modules[idx]:
                    fields = (i, syscall_name, syscall_addr, syscall_modname)
                    if hooked_tables.has_key(table):
                        hooked_tables[table].append(fields)
                    else:
                        hooked_tables[table] = [(fields)]

        return hooked_tables

    def calculate(self):

        if not has_distorm3:
            debug.warning("For best results please install distorm3")

        # Checks that subclass AbstractThreadCheck
        checks = registry.get_plugin_classes(AbstractThreadCheck)

        # If --listtags is chosen, just print the tags and return 
        if self._config.LISTTAGS:
            for cls_name, cls in checks.items():
                sys.stdout.write("{0:<20} {1}\n".format(cls_name, pydoc.getdoc(cls)))
            return

        addr_space = utils.load_as(self._config)
        system_range = tasks.get_kdbg(addr_space).MmSystemRangeStart.dereference_as("Pointer")

        # Only show threads owned by particular processes
        if self._config.PID:
            pidlist = [int(p) for p in self._config.PID.split(',')]
        else:
            pidlist = []

        # Get sorted list of kernel modules 
        mods = dict((addr_space.address_mask(mod.DllBase), mod) for mod in modules.lsmod(addr_space))
        mod_addrs = sorted(mods.keys())

        # Gather processes 
        all_tasks = list(tasks.pslist(addr_space))

        # Are we on x86 or x64. Save this for render_text 
        self.bits32 = addr_space.profile.metadata.\
            get("memory_model", "32bit") == "32bit"

        # Get a list of hooked SSDTs but only on x86
        if self.bits32:
            hooked_tables = self.get_hooked_tables(addr_space)
        else:
            hooked_tables = None

        # Dictionary to store threads. Keys are physical offsets of
        # ETHREAD objects. Values are tuples, where the first item is
        # a boolean specifying if the object was found by scanning and
        # the second item is the actual ETHREAD object. 
        seen_threads = dict()

        # Gather threads by list traversal of active/linked processes 
        for task in self.filter_tasks(all_tasks):
            for thread in task.ThreadListHead.\
                    list_of_type("_ETHREAD", "ThreadListEntry"):
                seen_threads[thread.obj_vm.vtop(thread.obj_offset)] = (False, thread)

        # Now scan for threads and save any that haven't been seen
        for thread in modscan.ThrdScan(self._config).calculate():
            if not seen_threads.has_key(thread.obj_offset):
                seen_threads[thread.obj_offset] = (True, thread)

        # Keep a record of processes whose DLLs we've already enumerated
        process_dll_info = {}

        for _offset, (found_by_scanner, thread) in seen_threads.items():

            # Skip processes the user doesn't want to see
            if pidlist and thread.Cid.UniqueProcess not in pidlist:
                continue

            # Do we need to gather DLLs for module resolution 
            if addr_space.address_compare(thread.StartAddress, system_range) != -1:
                owner = tasks.find_module(mods, 
                                          mod_addrs, 
                                          addr_space.address_mask(thread.StartAddress))
            else:
                owning_process = thread.owning_process() 
                if not owning_process.is_valid(): 
                    owner = None
                else:
                    try:
                        user_mod_addrs, user_mods = process_dll_info[owning_process.obj_offset]
                    except KeyError:
                        user_mods = dict((addr_space.address_mask(mod.DllBase), mod) 
                                            for mod in owning_process.get_load_modules())
                        user_mod_addrs = sorted(user_mods.keys())
                        process_dll_info[owning_process.obj_offset] = (user_mod_addrs, user_mods)
                    owner = tasks.find_module(user_mods, 
                                              user_mod_addrs, 
                                              addr_space.address_mask(thread.StartAddress))
            
            if owner:
                owner_name = str(owner.BaseDllName or '')
            else:
                owner_name = "UNKNOWN"

            # Replace the dummy class with an instance 
            instances = dict(
                        (cls_name, cls(thread, mods, mod_addrs,
                            hooked_tables, found_by_scanner))
                        for cls_name, cls in checks.items()
                        )

            yield thread, addr_space, mods, mod_addrs, \
                        instances, hooked_tables, system_range, owner_name

    def render_text(self, outfd, data):

        # Determine which filters the user wants to see
        if self._config.FILTER:
            filters = set(self._config.FILTER.split(','))
        else:
            filters = set()

        for thread, addr_space, mods, mod_addrs, \
                     instances, hooked_tables, system_range, owner_name in data:
            # If the user didn't set filters, display all results. If 
            # the user set one or more filters, only show threads 
            # with matching results. 
            tags = set([t for t, v in instances.items() if v.check()])

            if filters and not filters & tags:
                continue

            s = "------\n"

            s += "ETHREAD: {0:#010x} Pid: {1} Tid: {2}\n".format(
                thread.obj_offset,
                thread.Cid.UniqueProcess, thread.Cid.UniqueThread)

            s += "Tags: {0}\n".format(','.join(tags))
            s += "Created: {0}\n".format(thread.CreateTime)
            s += "Exited: {0}\n".format(thread.ExitTime)

            s += "Owning Process: {0}\n".format(
                thread.owning_process().ImageFileName)

            s += "Attached Process: {0}\n".format(
                thread.attached_process().ImageFileName)

            # Lookup the thread's state
            state = str(thread.Tcb.State)

            # Append the wait reason 
            if state == 'Waiting':
                state = state + ':' + str(thread.Tcb.WaitReason)

            s += "State: {0}\n".format(state)
            s += "BasePriority: {0:#x}\n".format(thread.Tcb.BasePriority)
            s += "Priority: {0:#x}\n".format(thread.Tcb.Priority)
            s += "TEB: {0:#010x}\n".format(thread.Tcb.Teb)

            s += "StartAddress: {0:#010x} {1}\n".format(
                thread.StartAddress, owner_name)

            # Check the flag which indicates whether Win32StartAddress is valid
            if thread.SameThreadApcFlags & 1:
                s += "Win32StartAddress: {0:#010x}\n".format(
                    thread.Win32StartAddress)

            if self.bits32:
                s += "ServiceTable: {0:#010x}\n".format(thread.Tcb.ServiceTable)

                ssdt_obj = obj.Object("_SERVICE_DESCRIPTOR_TABLE",
                    offset = thread.Tcb.ServiceTable,
                    vm = addr_space
                    )

                if ssdt_obj != None:
                    for i, desc in enumerate(ssdt_obj.Descriptors):
                        if desc.is_valid():
                            s += "  [{0}] {1:#010x}\n".format(i, desc.KiServiceTable.v())
                        else:
                            s += "  [{0}] -\n".format(i)
                        # Show exactly which functions are hooked 
                        table = desc.KiServiceTable.v()
                        if table not in hooked_tables.keys():
                            continue
                        for (i, func_name, func_addr, mod_name) in hooked_tables[table]:
                            s += "      [{0:#x}] {1} {2:#x} {3}\n".format(
                                i, func_name, func_addr, mod_name)

            s += "Win32Thread: {0:#010x}\n".format(thread.Tcb.Win32Thread)
            s += "CrossThreadFlags: {0}\n".format(thread.CrossThreadFlags)

            # Print the registers if possible 
            trapframe = thread.Tcb.TrapFrame.dereference_as("_KTRAP_FRAME")

            if trapframe and self.bits32:
                s += "Eip: {0:#10x}\n".format(trapframe.Eip)
                s += "  eax={0:#010x} ebx={1:#010x} ecx={2:#010x}".format(
                    trapframe.Eax, trapframe.Ebx, trapframe.Ecx)
                s += " edx={0:#010x} esi={1:#010x} edi={2:#010x}\n".format(
                    trapframe.Edx, trapframe.Esi, trapframe.Edi)
                s += "  eip={0:#010x} esp={1:#010x} ebp={2:#010x} err={3:#010x}\n".format(
                    trapframe.Eip, trapframe.HardwareEsp, trapframe.Ebp, trapframe.ErrCode)
                s += "  cs={0:#04x} ss={1:#04x} ds={2:#04x}".format(
                    trapframe.SegCs, trapframe.HardwareSegSs, trapframe.SegDs)
                s += " es={0:#04x} gs={1:#04x} fs={2:#04x} efl={3:#010x}\n".format(
                    trapframe.SegEs, trapframe.SegGs, trapframe.SegFs, trapframe.EFlags)
                s += "  dr0={0:#010x} dr1={1:#010x} dr2={2:#010x}".format(
                    trapframe.Dr0, trapframe.Dr1, trapframe.Dr2)
                s += " dr3={0:#010x} dr6={1:#010x} dr7={2:#010x}\n".format(
                    trapframe.Dr3, trapframe.Dr6, trapframe.Dr7)

            # Disasemble the start address if possible 
            process_space = thread.owning_process().get_process_address_space()
            
            if process_space.is_valid_address(thread.StartAddress):
                buf = process_space.zread(thread.StartAddress, 24)

                mode = "32bit" if self.bits32 else "64bit"

                s += "\n".join(["{0:#x} {1:<16} {2}".format(o, h, i)
                    for o, i, h in malfind.Disassemble(buf, thread.StartAddress.v(), mode)])

            outfd.write("{0}\n".format(s))
