# 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.win32.tasks as tasks
import volatility.plugins.modscan as modscan
import volatility.plugins.filescan as filescan
import volatility.plugins.overlays.windows.windows as windows

#--------------------------------------------------------------------------------
# object classes 
#--------------------------------------------------------------------------------

class _PSP_CID_TABLE(windows._HANDLE_TABLE): #pylint: disable-msg=W0212
    """Subclass the Windows handle table object for parsing PspCidTable"""

    def get_item(self, entry, handle_value = 0):

        p = obj.Object("address", entry.Object.v(), self.obj_vm)

        handle = obj.Object("_OBJECT_HEADER",
                offset = (p & ~7) -
                self.obj_vm.profile.get_obj_offset('_OBJECT_HEADER', 'Body'),
                vm = self.obj_vm)

        return handle

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

class MalwarePspCid(obj.ProfileModification):
    before = ['WindowsOverlay', 'WindowsVTypes']
    conditions = {'os': lambda x: x == 'windows'}

    def modification(self, profile):
        profile.vtypes.update({"_PSP_CID_TABLE" : profile.vtypes["_HANDLE_TABLE"]})
        profile.merge_overlay({"_KDDEBUGGER_DATA64" : [None,
            {'PspCidTable': [None,
                    ["pointer", ["pointer", ['_PSP_CID_TABLE']]]],
            }]})
        profile.object_classes.update({
            '_PSP_CID_TABLE': _PSP_CID_TABLE,
        })

#--------------------------------------------------------------------------------
# psxview plugin
#--------------------------------------------------------------------------------

class PsXview(common.AbstractWindowsCommand):
    "Find hidden processes with various process listings"

    def __init__(self, config, *args):
        common.AbstractWindowsCommand.__init__(self, config, *args)
        config.add_option("PHYSICAL-OFFSET", short_option = 'P', default = False,
                          help = "Physcal Offset", action = "store_true")

    def check_pslist(self, all_tasks):
        """Enumerate processes from PsActiveProcessHead"""
        return dict((p.obj_vm.vtop(p.obj_offset), p) for p in all_tasks)

    def check_psscan(self):
        """Enumerate processes with pool tag scanning"""
        return dict((p.obj_offset, p)
                    for p in filescan.PSScan(self._config).calculate())

    def check_thrdproc(self, _addr_space):
        """Enumerate processes indirectly by ETHREAD scanning"""
        ret = dict()

        for ethread in modscan.ThrdScan(self._config).calculate():
            if ethread.ExitTime != 0:
                continue
            # Bounce back to the threads owner 
            process = None
            if hasattr(ethread.Tcb, 'Process'):
                process = ethread.Tcb.Process.dereference_as('_EPROCESS')
            elif hasattr(ethread, 'ThreadsProcess'):
                process = ethread.ThreadsProcess.dereference()
            # Make sure the bounce succeeded 
            if (process and process.ExitTime == 0 and
                    process.UniqueProcessId > 0 and
                    process.UniqueProcessId < 65535):
                ret[process.obj_vm.vtop(process.obj_offset)] = process

        return ret

    def check_pspcid(self, addr_space):
        """Enumerate processes by walking the PspCidTable"""
        ret = dict()

        # Follow the pointers to the table base
        kdbg = tasks.get_kdbg(addr_space)
        PspCidTable = kdbg.PspCidTable.dereference().dereference()

        # Walk the handle table
        for handle in PspCidTable.handles():
            if handle.get_object_type() == "Process":
                process = handle.dereference_as("_EPROCESS")
                ret[process.obj_vm.vtop(process.obj_offset)] = process

        return ret

    def check_csrss_handles(self, all_tasks):
        """Enumerate processes using the csrss.exe handle table"""
        ret = dict()

        for p in all_tasks:
            if str(p.ImageFileName).lower() == "csrss.exe":
                # Gather the handles to process objects
                for handle in p.ObjectTable.handles():
                    if handle.get_object_type() == "Process":
                        process = handle.dereference_as("_EPROCESS")
                        ret[process.obj_vm.vtop(process.obj_offset)] = process

        return ret

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

        all_tasks = list(tasks.pslist(addr_space))

        ps_sources = {}
        # The keys are names of process sources. The values
        # are dictionaries whose keys are physical process 
        # offsets and the values are _EPROCESS objects. 
        ps_sources['pslist'] = self.check_pslist(all_tasks)
        ps_sources['psscan'] = self.check_psscan()
        ps_sources['thrdproc'] = self.check_thrdproc(addr_space)
        ps_sources['csrss'] = self.check_csrss_handles(all_tasks)
        ps_sources['pspcid'] = self.check_pspcid(addr_space)

        # Build a list of offsets from all sources
        seen_offsets = []
        for source in ps_sources.values():
            for offset in source.keys():
                if offset not in seen_offsets:
                    seen_offsets.append(offset)
                    yield offset, source[offset], ps_sources

    def render_text(self, outfd, data):

        self.table_header(outfd, [('Offset(P)', '[addrpad]'),
                                  ('Name', '<20'),
                                  ('PID', '>6'),
                                  ('pslist', '5'),
                                  ('psscan', '5'),
                                  ('thrdproc', '5'),
                                  ('pspcdid', '5'),
                                  ('csrss', '5')
                                  ])

        for offset, process, ps_sources in data:
            self.table_row(outfd,
                offset,
                process.ImageFileName,
                process.UniqueProcessId,
                str(ps_sources['pslist'].has_key(offset)),
                str(ps_sources['psscan'].has_key(offset)),
                str(ps_sources['thrdproc'].has_key(offset)),
                str(ps_sources['pspcid'].has_key(offset)),
                str(ps_sources['csrss'].has_key(offset))
                )
