# 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
#

# pylint: disable=protected-access


from rekall import utils
from rekall.plugins.windows import common

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

    __name = "psxview"

    def __init__(self, **kwargs):
        """Find hidden processes with various process listings methods."""
        super(PsXview, self).__init__(**kwargs)

    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"""
        result = {}
        psscan = self.session.plugins.psscan(session=self.session)
        for hit in psscan.scan_processes():
            result[hit.obj_offset] = hit

        return result

    def check_thrdproc(self):
        """Enumerate processes indirectly by ETHREAD scanning"""
        thrdscan_plugin = self.session.plugins.thrdscan()
        for ethread in thrdscan_plugin.generate_hits():
            if ethread.ExitTime != 0:
                continue

            # Bounce back to the threads owner
            process = ethread.Tcb.m('Process').dereference_as(
                    '_EPROCESS', vm=self.kernel_address_space)

            if not process:
                process = ethread.m('ThreadsProcess').dereference(
                    vm=self.kernel_address_space)

            # Make sure the bounce succeeded
            if (process and process.ExitTime == 0 and
                    process.UniqueProcessId > 0 and
                    process.UniqueProcessId < 0xFFFF):

                yield process.obj_vm.vtop(process.obj_offset)

    def check_pspcid(self):
        """Enumerate processes by walking the PspCidTable"""
        # Follow the pointers to the table base
        PspCidTable = self.kdbg.PspCidTable.cast("Pointer",
                                                 target="_PSP_CID_TABLE")

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

    def check_csrss_handles(self):
        """Enumerate processes using the csrss.exe handle table"""
        for task in self.session.plugins.pslist().list_eprocess():
            if utils.SmartUnicode(task.ImageFileName).lower() == "csrss.exe":
                # Gather the handles to process objects
                for handle in task.ObjectTable.handles():
                    if handle.get_object_type() == "Process":
                        process = handle.dereference_as("_EPROCESS")
                        yield process.obj_vm.vtop(process.obj_offset)

    def find_processes(self):
        ps_sources = {}

        pslist_plugin = self.session.plugins.pslist()
        # We store the physical offsets of the _EPROCESS structs.
        ps_sources['pslist'] = set([x.obj_vm.vtop(x.obj_offset)
                                    for x in pslist_plugin.list_eprocess()])

        psscan_plugin = self.session.plugins.psscan(
            address_space=self.physical_address_space)

        ps_sources['psscan'] = set([x.obj_offset
                                    for x in psscan_plugin.scan_processes()])
        ps_sources['thrdproc'] = set(self.check_thrdproc())
        ps_sources['csrss'] = set(self.check_csrss_handles())
        ps_sources['pspcid'] = set(self.check_pspcid())

        # Build a list of offsets from all sources
        all_offsets = set().union(*ps_sources.values())

        for offset in all_offsets:
            # These are the sets that contain this offset
            contains = [k for k, v in ps_sources.items() if offset in v]

            phys_eprocess = self.profile._EPROCESS(
                vm=self.physical_address_space,
                offset=offset)

            virtual_eprocess = psscan_plugin.guess_eprocess_virtual_address(
                phys_eprocess)

            yield phys_eprocess, virtual_eprocess, contains

    def render(self, renderer):
        renderer.table_header(
            [('Offset(P)', 'physical_offset', '[addrpad]'),
             ('Offset(V)', 'virtual_offset', '[addrpad]'),
             ('Name', 'name', '<20'),
             ('PID', 'pid', '>6'),
             ('pslist', 'pslist', '7'),
             ('psscan', 'psscan', '7'),
             ('thrdproc', 'thrdproc', '7'),
             ('pspcdid', 'pspcdid', '7'),
             ('csrss', 'csrss', '7')
             ])

        for process, virtual_eprocess, contains in self.find_processes():
            renderer.table_row(
                process,
                virtual_eprocess,
                process.ImageFileName,
                process.UniqueProcessId,
                ('pslist' in contains),
                ('psscan' in contains),
                ('thrdproc' in contains),
                ('pspcid' in contains),
                ('csrss' in contains),
                )
