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

# File-wide pylint filter for protected members, since we have three _BLAH structures
#pylint: disable-msg=W0212

import os
import volatility.utils as utils
import volatility.obj as obj
import volatility.debug as debug
import volatility.win32.tasks as tasks
import volatility.win32.modules as modules
import volatility.plugins.taskmods as taskmods
import volatility.plugins.vadinfo as vadinfo
import volatility.plugins.overlays.windows.windows as windows
import volatility.constants as constants

try:
    import yara
    has_yara = True
except ImportError:
    has_yara = False

try:
    import distorm3
    has_distorm3 = True
except ImportError:
    has_distorm3 = False

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

class MalwareEPROCESS(windows._EPROCESS):
    """Extension of the default EPROCESS with some helpers"""

    @property
    def IsWow64(self):
        """Returns True if this is a wow64 process"""
        return hasattr(self, 'Wow64Process') and self.Wow64Process.v() != 0

    @property
    def SessionId(self):
        """Returns the Session ID of the process"""

        if self.Session.is_valid():
            process_space = self.get_process_address_space()
            if process_space:
                return obj.Object("_MM_SESSION_SPACE",
                                  offset = self.Session,
                                  vm = process_space).SessionId

        return obj.NoneObject("Cannot find process session")

    def get_vads(self, vad_filter = None):
        """
        Generator for MMVADs that match specific
        metadata.

        @param vad_filter: a callable that is passed the 
        current MMVAD and applies tests to the MMVAD struct
        members or nested struct members. 

        @yields a tuple (mmvad, address_space). Where mmvad is
        the MMVAD object in kernel AS and address_space 
        is the process address space. 
        """

        # We absolutely need a process AS. If this 
        # fails then all else fails
        process_space = self.get_process_address_space()
        if not process_space:
            return

        for vad in self.VadRoot.traverse():
            if not vad.is_valid():
                continue
            # Apply the meta filter if one is supplied
            if vad_filter:
                if not vad_filter(vad):
                    continue
            yield vad, process_space

    def search_process_memory(self, s):
        """
        Search memory for a simple byte string.

        @param s: the string to search for.

        @returns every occurrance of the string 
        in process memory (as absolute address).
        """

        def iterfind(data, string):
            offset = data.find(string, 0)
            while offset >= 0:
                yield offset
                offset = data.find(string, offset + len(string))

        # Allow for some overlap in case objects are 
        # right on page boundaries 
        overlap = 1024

        # All MMVADs that belong to this process.
        for vad, address_space in self.get_vads():
            offset = vad.Start
            out_of_range = vad.Start + vad.Length
            while offset < out_of_range:
                # Read some data and match it.
                to_read = min(constants.SCAN_BLOCKSIZE + overlap, out_of_range - offset)
                data = address_space.zread(offset, to_read)
                if not data:
                    break
                for hit in iterfind(data, s):
                    yield offset + hit
                offset += min(to_read, constants.SCAN_BLOCKSIZE)

    def _injection_filter(self, vad):
        """
        This is a callback that's executed by get_vads()
        when searching for injected code / hidden DLLs. 

        This looks for private allocations that are committed, 
        memory-resident, non-empty (not all zeros) and with an 
        original protection that includes write and execute. 

        It is important to note that protections are applied at 
        the allocation granularity (page level). Thus the original
        protection might not be the current protection, and it
        also might not apply to all pages in the VAD range. 

        @param vad: an MMVAD object.

        @returns: True if the MMVAD looks like it might
        contain injected code. 
        """
        protect = vadinfo.PROTECT_FLAGS.get(vad.u.VadFlags.Protection.v(), "")
        write_exec = "EXECUTE" in protect and "WRITE" in protect

        # The Write/Execute check applies to everything 
        if not write_exec:
            return False

        # This is a typical VirtualAlloc'd injection 
        if vad.u.VadFlags.PrivateMemory == 1 and vad.Tag == "VadS":
            return True

        # This is a stuxnet-style injection 
        if (vad.u.VadFlags.PrivateMemory == 0 and
                protect != "PAGE_EXECUTE_WRITECOPY"):
            return True

        return False

    def _mapped_file_filter(self, vad):
        """
        This is a callback that's executed by get_vads() 
        when searching for memory-mapped files. 

        @param vad: an MMVAD object.

        @returns: True if the MMVAD looks like it might
        contain a mapped file. 
        """

        return vad.u.VadFlags.PrivateMemory == 0 and vad.ControlArea

    def environment_variables(self):
        """Generator for environment variables. 

        The PEB points to our env block - a series of null-terminated
        unicode strings. Each string cannot be more than 0x7FFF chars. 
        End of the list is a quad-null. 
        """

        # Address of the environment block
        if not self.Peb.ProcessParameters.Environment.is_valid():
            return

        process_space = self.get_process_address_space()
        if not process_space:
            return

        block = self.Peb.ProcessParameters.Environment

        s = obj.Object("String", offset = block, vm = process_space,
            encoding = 'utf16', length = 0x7FFF)

        # The terminator is a quad null 
        while len(s):
            if s.count(u"=") == 1:
                yield s.split(u"=")
            # Scan forward the length of this string plus the null
            next_offset = s.obj_offset + ((len(s) + 1) * 2)
            s = obj.Object("String", offset = next_offset,
                vm = process_space, encoding = 'utf16', length = 0x7FFF)

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

class MalwareObjectClasesXP(obj.ProfileModification):
    before = ['WindowsObjectClasses']
    conditions = {'os': lambda x: x == 'windows'}
    def modification(self, profile):
        profile.object_classes.update({
            '_EPROCESS': MalwareEPROCESS,
        })

#--------------------------------------------------------------------------------
# functions 
#--------------------------------------------------------------------------------

def Disassemble(data, start, bits = '32bit', stoponret = False):
    """Dissassemble code with distorm3. 

    @param data: python byte str to decode
    @param start: address where `data` is found in memory
    @param bits: use 32bit or 64bit decoding 
    @param stoponret: stop disasm when function end is reached
    
    @returns: tuple of (offset, instruction, hex bytes)
    """

    if not has_distorm3:
        raise StopIteration

    if bits == '32bit':
        mode = distorm3.Decode32Bits
    else:
        mode = distorm3.Decode64Bits

    for o, _, i, h in distorm3.DecodeGenerator(start, data, mode):
        if stoponret and i.startswith("RET"):
            raise StopIteration
        yield o, i, h

#--------------------------------------------------------------------------------
# scanners by scudette
#
# unfortunately the existing scanning framework (i.e. scan.BaseScanner) has 
# some shortcomings that don't allow us to integrate yara easily. 
#
# FIXME: these may need updating after resolving issue 310 which aims to 
# enhance the scan.BaseScanner to better support things like this
#--------------------------------------------------------------------------------

class BaseYaraScanner(object):
    """An address space scanner for Yara signatures."""
    overlap = 1024

    def __init__(self, address_space = None, rules = None):
        self.rules = rules
        self.address_space = address_space

    def scan(self, offset, maxlen):
        # Start scanning from offset until maxlen:
        i = offset

        while i < offset + maxlen:
            # Read some data and match it.
            to_read = min(constants.SCAN_BLOCKSIZE + self.overlap, offset + maxlen - i)
            data = self.address_space.zread(i, to_read)
            if data:
                for match in self.rules.match(data = data):
                    # We currently don't use name or value from the 
                    # yara results but they can be yielded in the 
                    # future if necessary. 
                    for moffset, _name, _value in match.strings:
                        if moffset < constants.SCAN_BLOCKSIZE:
                            yield match, moffset + i

            i += constants.SCAN_BLOCKSIZE

class VadYaraScanner(BaseYaraScanner):
    """A scanner over all memory regions of a process."""

    def __init__(self, task = None, **kwargs):
        """Scan the process address space through the Vads.

        Args:
          task: The _EPROCESS object for this task.
        """
        self.task = task
        BaseYaraScanner.__init__(self, address_space = task.get_process_address_space(), **kwargs)

    def scan(self, offset = 0, maxlen = None):
        for vad, self.address_space in self.task.get_vads():
            for match in BaseYaraScanner.scan(self, vad.Start, vad.Length):
                yield match

class DiscontigYaraScanner(BaseYaraScanner):
    """A Scanner for Discontiguous scanning."""

    def scan(self, start_offset = 0, maxlen = None):
        contiguous_offset = 0
        total_length = 0
        for (offset, length) in self.address_space.get_available_addresses():
            # Skip ranges before the start_offset
            if offset < start_offset:
                continue

            # Try to join up adjacent pages as much as possible.
            if offset == contiguous_offset + total_length:
                total_length += length
            else:
                # Scan the last contiguous range.
                for match in BaseYaraScanner.scan(self, contiguous_offset, total_length):
                    yield match

                # Reset the contiguous range.
                contiguous_offset = offset
                total_length = length

        if total_length > 0:
            # Do the last range.
            for match in BaseYaraScanner.scan(self, contiguous_offset, total_length):
                yield match

#--------------------------------------------------------------------------------
# yarascan
#--------------------------------------------------------------------------------

class YaraScan(taskmods.DllList):
    "Scan process or kernel memory with Yara signatures"

    def __init__(self, config, *args, **kwargs):
        taskmods.DllList.__init__(self, config, *args, **kwargs)
        config.add_option("KERNEL", short_option = 'K', default = False, action = 'store_true',
                        help = 'Scan kernel modules')
        config.add_option("WIDE", short_option = 'W', default = False, action = 'store_true',
                        help = 'Match wide (unicode) strings')
        config.add_option('YARA-RULES', short_option = 'Y', default = None,
                        help = 'Yara rules (as a string)')
        config.add_option('YARA-FILE', short_option = 'y', default = None,
                        help = 'Yara rules (rules file)')
        config.add_option('DUMP-DIR', short_option = 'D', default = None,
                        help = 'Directory in which to dump the files')

    def calculate(self):

        if not has_yara:
            debug.error("Please install Yara from code.google.com/p/yara-project")

        addr_space = utils.load_as(self._config)

        if self._config.YARA_RULES:
            s = self._config.YARA_RULES
            # Don't wrap hex or regex rules in quotes 
            if s[0] not in ("{", "/"): s = '"' + s + '"'
            # Scan for unicode strings 
            if self._config.WIDE: s += "wide"
            rules = yara.compile(sources = {
                        'n' : 'rule r1 {strings: $a = ' + s + ' condition: $a}'
                        })
        elif self._config.YARA_FILE:
            rules = yara.compile(self._config.YARA_FILE)
        else:
            debug.error("You must specify a string (-Y) or a rules file (-y)")

        if self._config.KERNEL:

            # Find KDBG so we know where kernel memory begins. Do not assume
            # the starting range is 0x80000000 because we may be dealing with
            # an image with the /3GB boot switch. 
            kdbg = tasks.get_kdbg(addr_space)

            start = kdbg.MmSystemRangeStart.dereference_as("Pointer")

            # Modules so we can map addresses to owners
            mods = dict((mod.DllBase, mod)
                        for mod in modules.lsmod(addr_space))
            mod_addrs = sorted(mods.keys())

            # There are multiple views (GUI sessions) of kernel memory.
            # Since we're scanning virtual memory and not physical, 
            # all sessions must be scanned for full coverage. This 
            # really only has a positive effect if the data you're
            # searching for is in GUI memory. 
            sessions = []

            for proc in tasks.pslist(addr_space):
                sid = proc.SessionId
                # Skip sessions we've already seen 
                if sid == None or sid in sessions:
                    continue

                session_space = proc.get_process_address_space()
                if session_space == None:
                    continue

                sessions.append(sid)
                scanner = DiscontigYaraScanner(address_space = session_space,
                                               rules = rules)

                for hit, address in scanner.scan(start_offset = start):
                    module = tasks.find_module(mods, mod_addrs, address)
                    yield (module, address, hit, session_space.zread(address, 1024))

        else:
            for task in self.filter_tasks(tasks.pslist(addr_space)):
                scanner = VadYaraScanner(task = task, rules = rules)
                for hit, address in scanner.scan():
                    yield (task, address, hit, scanner.address_space.zread(address, 1024))

    def render_text(self, outfd, data):

        if self._config.DUMP_DIR and not os.path.isdir(self._config.DUMP_DIR):
            debug.error(self._config.DUMP_DIR + " is not a directory")

        for o, addr, hit, content in data:
            outfd.write("Rule: {0}\n".format(hit.rule))

            # Find out if the hit is from user or kernel mode 
            if o == None:
                outfd.write("Owner: (Unknown Kernel Memory)\n")
                filename = "kernel.{0:#x}.dmp".format(addr)
            elif o.obj_name == "_EPROCESS":
                outfd.write("Owner: Process {0} Pid {1}\n".format(o.ImageFileName,
                    o.UniqueProcessId))
                filename = "process.{0:#x}.{1:#x}.dmp".format(o.obj_offset, addr)
            else:
                outfd.write("Owner: {0}\n".format(o.BaseDllName))
                filename = "kernel.{0:#x}.{1:#x}.dmp".format(o.obj_offset, addr)

            # Dump the data if --dump-dir was supplied
            if self._config.DUMP_DIR:
                path = os.path.join(self._config.DUMP_DIR, filename)
                fh = open(path, "wb")
                fh.write(content)
                fh.close()

            outfd.write("".join(
                ["{0:#010x}  {1:<48}  {2}\n".format(addr + o, h, ''.join(c))
                for o, h, c in utils.Hexdump(content[0:64])
                ]))

#--------------------------------------------------------------------------------
# malfind
#--------------------------------------------------------------------------------

class Malfind(vadinfo.VADDump):
    "Find hidden and injected code"

    def _is_vad_empty(self, vad, address_space):
        """
        Check if a VAD region is either entirely unavailable 
        due to paging, entirely consiting of zeros, or a 
        combination of the two. This helps ignore false positives
        whose VAD flags match task._injection_filter requirements
        but there's no data and thus not worth reporting it. 

        @param vad: an MMVAD object in kernel AS
        @param address_space: the process address space 
        """
        
        PAGE_SIZE = 0x1000
        all_zero_page = "\x00" * PAGE_SIZE

        offset = 0
        while offset < vad.Length:
            next_addr = vad.Start + offset
            if (address_space.is_valid_address(next_addr) and 
                    address_space.read(next_addr, PAGE_SIZE) != all_zero_page):
                return False
            offset += PAGE_SIZE

        return True

    def render_text(self, outfd, data):

        if self._config.DUMP_DIR and not os.path.isdir(self._config.DUMP_DIR):
            debug.error(self._config.DUMP_DIR + " is not a directory")

        for task in data:
            for vad, address_space in task.get_vads(vad_filter = task._injection_filter):

                if self._is_vad_empty(vad, address_space):
                    continue

                content = address_space.zread(vad.Start, 64)

                outfd.write("Process: {0} Pid: {1} Address: {2:#x}\n".format(
                    task.ImageFileName, task.UniqueProcessId, vad.Start))

                outfd.write("Vad Tag: {0} Protection: {1}\n".format(
                    vad.Tag, vadinfo.PROTECT_FLAGS.get(vad.u.VadFlags.Protection.v(), "")))

                outfd.write("Flags: {0}\n".format(str(vad.u.VadFlags)))
                outfd.write("\n")

                outfd.write("{0}\n".format("\n".join(
                    ["{0:#010x}  {1:<48}  {2}".format(vad.Start + o, h, ''.join(c))
                    for o, h, c in utils.Hexdump(content)
                    ])))

                outfd.write("\n")
                outfd.write("\n".join(
                    ["{0:#x} {1:<16} {2}".format(o, h, i)
                    for o, i, h in Disassemble(content, vad.Start)
                    ]))

                # Dump the data if --dump-dir was supplied
                if self._config.DUMP_DIR:

                    filename = os.path.join(self._config.DUMP_DIR,
                        "process.{0:#x}.{1:#x}.dmp".format(
                        task.obj_offset, vad.Start))

                    self.dump_vad(filename, vad, address_space)

                outfd.write("\n\n")

#--------------------------------------------------------------------------------
# ldrmodules 
#--------------------------------------------------------------------------------

class LdrModules(taskmods.DllList):
    "Detect unlinked DLLs"

    def render_text(self, outfd, data):

        self.table_header(outfd,
            [("Pid", "8"),
             ("Process", "20"),
             ("Base", "[addrpad]"),
             ("InLoad", "5"),
             ("InInit", "5"),
             ("InMem", "5"),
             ("MappedPath", "")
            ])

        for task in data:
            # Build a dictionary for all three PEB lists where the
            # keys are base address and module objects are the values
            inloadorder = dict((mod.DllBase.v(), mod)
                                for mod in task.get_load_modules())
            ininitorder = dict((mod.DllBase.v(), mod)
                                for mod in task.get_init_modules())
            inmemorder = dict((mod.DllBase.v(), mod)
                                for mod in task.get_mem_modules())

            # Build a similar dictionary for the mapped files 
            mapped_files = {}
            for vad, address_space in task.get_vads(vad_filter = task._mapped_file_filter):
                # Note this is a lot faster than acquiring the full
                # vad region and then checking the first two bytes. 
                if obj.Object("_IMAGE_DOS_HEADER", offset = vad.Start, vm = address_space).e_magic != 0x5A4D:
                    continue
                mapped_files[int(vad.Start)] = vad.FileObject.FileName

            # For each base address with a mapped file, print info on 
            # the other PEB lists to spot discrepancies. 
            for base in mapped_files.keys():
                # Does the base address exist in the PEB DLL lists?
                load_mod = inloadorder.get(base, None)
                init_mod = ininitorder.get(base, None)
                mem_mod = inmemorder.get(base, None)
                # Report if the mapped files are in the PEB lists
                self.table_row(outfd,
                        task.UniqueProcessId,
                        task.ImageFileName,
                        base,
                        str(load_mod != None),
                        str(init_mod != None),
                        str(mem_mod != None),
                        mapped_files[base]
                        )
                # Print the full paths and base names in verbose mode 
                if self._config.verbose:
                    if load_mod:
                        outfd.write("  Load Path: {0} : {1}\n".format(load_mod.FullDllName, load_mod.BaseDllName))
                    if init_mod:
                        outfd.write("  Init Path: {0} : {1}\n".format(init_mod.FullDllName, init_mod.BaseDllName))
                    if mem_mod:
                        outfd.write("  Mem Path:  {0} : {1}\n".format(mem_mod.FullDllName, mem_mod.BaseDllName))
