# Rekall Memory Forensics
# Copyright (c) 2012, Michael Cohen <scudette@gmail.com>
# 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
#

import yara

from rekall import plugin
from rekall import utils
from rekall.plugins import yarascanner
from rekall.plugins.windows import common
from rekall.plugins.windows import vadinfo


class VadYaraScanner(vadinfo.VadScanner, yarascanner.BaseYaraASScanner):
    """A Yara scanner which only operates on VAD regions."""


class YaraScan(common.WinProcessFilter):
    """Scan using yara signatures."""

    name = "yarascan"

    @classmethod
    def args(cls, parser):
        super(YaraScan, cls).args(parser)
        parser.add_argument("--string", default=None,
                            help="A verbatim string to search for.")

        parser.add_argument("--binary_string", default=None,
                            help="A binary string (encoded as hex) to search "
                            "for. e.g. 000102[1-200]0506")

        parser.add_argument("--yara_file", default=None,
                            help="The yara signature file to read.")

        parser.add_argument("--yara_expression", default=None,
                            help="If provided we scan for this yarra "
                            "expression.")

        parser.add_argument(
            "--scan_vads", default=False, type="Boolean",
            help=("If specified we scan the vads of "
                  "processes, else we scan their entire address spaces. "
                  "Note that scanning the entire address space will "
                  "typically also include the kernel, while targetting "
                  "VADs only will only include the addresses mapped by "
                  "the process."))

        parser.add_argument(
            "--scan_physical", default=False, type="Boolean",
            help="If specified we scan the physcial address space.")

    def __init__(self, string=None, scan_vads=False, scan_physical=False,
                 yara_file=None, yara_expression=None, binary_string=None,
                 **kwargs):
        """Scan using yara signatures."""
        super(YaraScan, self).__init__(**kwargs)
        if yara_expression:
            self.rules_source = yara_expression
            self.rules = yara.compile(source=self.rules_source)

        elif binary_string:
            self.compile_rule(
                'rule r1 {strings: $a = {%s} condition: $a}' % binary_string
                )
        elif string:
            self.compile_rule(
                'rule r1 {strings: $a = "%s" condition: $a}' % string
                )

        elif yara_file:
            self.compile_rule(open(yara_file).read())
        else:
            raise plugin.PluginError("You must specify a yara rule file or "
                                     "string to match.")

        self.scan_vads = scan_vads
        self.scan_physical = scan_physical

    def compile_rule(self, rule):
        self.rules_source = rule
        try:
            self.rules = yara.compile(source=rule)
        except Exception as e:
            raise plugin.PluginError(
                "Failed to compile yara expression: %s" % e)

    def generate_hits(self, address_space):
        scanner = yarascanner.BaseYaraASScanner(
            profile=self.profile, session=self.session,
            address_space=address_space,
            rules=self.rules)

        return scanner.scan()

    def render_scan_physical(self, renderer):
        """This method scans the process memory using the VAD."""
        for rule, address, _, hit in self.generate_hits(
            self.physical_address_space):
            renderer.format("Rule: {0}\n", rule)

            context = self.physical_address_space.read(address, 0x40)
            utils.WriteHexdump(renderer, context, base=address)

    def render_kernel_scan(self, renderer):
        modules = self.session.plugins.modules()

        for rule, address, _, hit in self.generate_hits(
            self.kernel_address_space):
            renderer.format("Rule: {0}\n", rule)

            # Find out who owns this hit.
            owner = modules.find_module(address)
            if owner:
                renderer.format("Owner: {0}\n", owner.BaseDllName)
                filename = "kernel.{0:#x}.{1:#x}.dmp".format(
                    owner.obj_offset, address)
            else:
                filename = "kernel.{0:#x}.dmp".format(
                    address)

                renderer.format("Owner: (Unknown Kernel Memory)\n")

            context = self.kernel_address_space.read(address, 0x40)
            utils.WriteHexdump(renderer, context, base=address)

    def render_task_scan(self, renderer, task):
        task_as = task.get_process_address_space()

        for rule, address, _, hit in self.generate_hits(task_as):
            renderer.format("Rule: {0}\n", rule)

            renderer.format("Owner: {0}\n", task.ImageFileName)

            context = task_as.read(address, 0x40)
            utils.WriteHexdump(renderer, context, base=address)

    def render_task_scan_vad(self, renderer, task):
        """This method scans the process memory using the VAD."""
        task_as = task.get_process_address_space()

        scanner = VadYaraScanner(
            session=self.session, rules=self.rules, task=task)

        for rule, address, _, hit in scanner.scan():
            renderer.format("Rule: {0}\n", rule)

            renderer.format("Owner: {0}\n", task.ImageFileName)

            context = task_as.read(address, 0x40)
            utils.WriteHexdump(renderer, context, base=address)


    def render(self, renderer):
        """Render output."""
        if self.scan_physical:
            return self.render_scan_physical(renderer)

        elif self.scan_vads or self.filtering_requested:
            for task in self.filter_processes():
                if self.scan_vads:
                    self.render_task_scan_vad(renderer, task)
                else:
                    self.render_task_scan(renderer, task)

        # We are searching the kernel address space
        else:
            return self.render_kernel_scan(renderer)
