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

"""A Rekall Memory Forensics scanner which uses yara."""
import sys
import yara

from rekall import constants
from rekall import plugin
from rekall import scan
from rekall import utils
from rekall.plugins.windows import common
from rekall.plugins.windows import vadinfo


class BaseYaraASScanner(scan.BaseScanner):
    """An address space scanner for Yara signatures."""
    overlap = 1024

    def __init__(self, rules=None, **kwargs):
        super(BaseYaraASScanner, self).__init__(**kwargs)
        self.rules = rules
        self.hits = []
        self.base_offset = None

    def _match_rules(self, buffer_as):
        """Compatibility for yara modules.

        Unfortunately there are two different implementations of the yara python
        bindings:

        # The original upstream source.
        http://plusvic.github.io/yara/

        # The version which is installed using pip install.
        https://github.com/mjdorma/yara-ctypes

        These do not work the same and so we need to support both.

        Yields:
          a tuple of (offset, rule_name, name, value)
        """
        matches = self.rules.match(data=buffer_as.data)
        # yara-cpython bindings from pip.
        if type(matches) is dict:
            for source, matches in matches.items():
                for match in matches:
                    for string in match["strings"]:
                        hit_offset = string["offset"] + buffer_as.base_offset

                        yield (match["rule"], hit_offset,
                               string["identifier"], string["data"])

        else:
            # native bindings from http://plusvic.github.io/yara/
            for match in matches:
                for buffer_offset, name, value in match.strings:
                    hit_offset = buffer_offset + buffer_as.base_offset
                    yield (match.rule, hit_offset, name, value)

    def check_addr(self, offset, buffer_as=None):
        # The buffer was changed - we scan the entire buffer and record the
        # hits - then we can feed it to the Rekall scan framework.
        if self.base_offset != buffer_as.base_offset:
            self.base_offset = buffer_as.base_offset
            self.hits = []

            for rule, offset, name, value in self._match_rules(buffer_as):
                self.hits.append((rule, offset, name, value))

        if self.hits and offset == self.hits[0][1]:
            return self.hits.pop(0)

    def skip(self, buffer_as, offset):
        # Skip the entire buffer.
        if not self.hits:
            return len(buffer_as.data)

        next_hit = self.hits[0][1]
        return next_hit - offset


class VadYarraScanner(vadinfo.VadScanner, BaseYaraASScanner):
    """A Yarra 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("--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, action="store_true",
                            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, action="store_true",
                            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, **kwargs):
        """Scan using yara signatures.

        Args:
          string: A verbatim string to search for.
          scan_vads: If true we scan the vads of the specified processes, else
            we scan their entire address spaces.
          scan_physical: If true we scan the physical address space.
          yara_file: The yara file to read.
          yara_expression: If provided we scan for this yarra expression.
        """
        super(YaraScan, self).__init__(**kwargs)
        if yara_expression:
            self.rules_source = yara_expression
            self.rules = yara.compile(source=self.rules_source)

        elif string:
            self.rules_source = (
                'rule r1 {strings: $a = "%s" condition: $a}' % string
                )
            self.rules = yara.compile(source=self.rules_source)

        elif yara_file:
            self.rules = yara.compile(yara_file)
        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 generate_hits(self, address_space):
        scanner = 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 = VadYarraScanner(
            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)
