# 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.plugins.windows import common


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

    __name = "psxview"

    METHODS = common.WinProcessFilter.METHODS + [
        "PSScan", "Thrdproc"]

    __args = [
        dict(name="method", choices=METHODS, type="ChoiceArray",
             default=METHODS, help="Method to list processes.",
             override=True),
    ]

    def render(self, renderer):
        headers = [
            dict(type="_EPROCESS", name="_EPROCESS"),
            ]

        for method in self.plugin_args.method:
            headers.append((method, method, "%s" % len(method)))

        renderer.table_header(headers)

        for eprocess in self.filter_processes():
            row = [eprocess]

            for method in self.plugin_args.method:
                row.append(eprocess.obj_offset in
                           self.session.GetParameter("pslist_%s" % method))
            renderer.table_row(*row)


class PsListPSScanHook(common.AbstractWindowsParameterHook):
    name = "pslist_PSScan"

    def calculate(self):
        """Enumerate processes by scanning for _EPROCESS."""
        result = set()

        psscan = self.session.plugins.psscan()
        pslist = self.session.plugins.pslist()
        for row in psscan.collect():
            physical_eprocess = row["offset_p"]
            if physical_eprocess.obj_vm == self.session.physical_address_space:
                eprocess = pslist.virtual_process_from_physical_offset(
                    physical_eprocess)
            else:
                eprocess = physical_eprocess

            if eprocess != None:
                result.add(eprocess.obj_offset)

        self.session.logging.debug(
            "Listed %s processes using PSScan", len(result))

        return result


class PsListThrdprocHook(common.AbstractWindowsParameterHook):
    name = "pslist_Thrdproc"

    def calculate(self):
        """Enumerate processes by scanning for threads."""
        result = set()

        thrdscan_plugin = self.session.plugins.thrdscan()
        for results in thrdscan_plugin.collect():
            ethread = self.session.profile._ETHREAD(results[0])

            if ethread.ExitTime != 0:
                continue

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

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

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

                result.add(process.obj_offset)

        self.session.logging.debug(
            "Listed %s processes using Thrdproc", len(result))

        return result
