# Vortessence  Memory Forensics
#
# Copyright 2007-2015 Volatility Foundation
# Copyright 2010, 2011, 2012 Michael Ligh <michael.ligh@mnin.org>
#
# Copyright 2015 Bern University of Applied Sciences
#
# 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 3 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, see <http://www.gnu.org/licenses/>.

import json
import re
import volatility.obj as obj
import volatility.plugins.filescan as filescan
import volatility.win32.modules as modules
import volatility.win32.tasks as tasks
import volatility.utils as utils
import volatility.plugins.malware.malfind as malfind
import volatility.plugins.overlays.windows.windows as windows

#--------------------------------------------------------------------------------
# constants
#--------------------------------------------------------------------------------

MAJOR_FUNCTIONS = [
    'IRP_MJ_CREATE',
    'IRP_MJ_CREATE_NAMED_PIPE',
    'IRP_MJ_CLOSE',
    'IRP_MJ_READ',
    'IRP_MJ_WRITE',
    'IRP_MJ_QUERY_INFORMATION',
    'IRP_MJ_SET_INFORMATION',
    'IRP_MJ_QUERY_EA',
    'IRP_MJ_SET_EA',
    'IRP_MJ_FLUSH_BUFFERS',
    'IRP_MJ_QUERY_VOLUME_INFORMATION',
    'IRP_MJ_SET_VOLUME_INFORMATION',
    'IRP_MJ_DIRECTORY_CONTROL',
    'IRP_MJ_FILE_SYSTEM_CONTROL',
    'IRP_MJ_DEVICE_CONTROL',
    'IRP_MJ_INTERNAL_DEVICE_CONTROL',
    'IRP_MJ_SHUTDOWN',
    'IRP_MJ_LOCK_CONTROL',
    'IRP_MJ_CLEANUP',
    'IRP_MJ_CREATE_MAILSLOT',
    'IRP_MJ_QUERY_SECURITY',
    'IRP_MJ_SET_SECURITY',
    'IRP_MJ_POWER',
    'IRP_MJ_SYSTEM_CONTROL',
    'IRP_MJ_DEVICE_CHANGE',
    'IRP_MJ_QUERY_QUOTA',
    'IRP_MJ_SET_QUOTA',
    'IRP_MJ_PNP'
]

DEVICE_CODES = {
    0x00000027 : 'FILE_DEVICE_8042_PORT',
    0x00000032 : 'FILE_DEVICE_ACPI',
    0x00000029 : 'FILE_DEVICE_BATTERY',
    0x00000001 : 'FILE_DEVICE_BEEP',
    0x0000002a : 'FILE_DEVICE_BUS_EXTENDER',
    0x00000002 : 'FILE_DEVICE_CD_ROM',
    0x00000003 : 'FILE_DEVICE_CD_ROM_FILE_SYSTEM',
    0x00000030 : 'FILE_DEVICE_CHANGER',
    0x00000004 : 'FILE_DEVICE_CONTROLLER',
    0x00000005 : 'FILE_DEVICE_DATALINK',
    0x00000006 : 'FILE_DEVICE_DFS',
    0x00000035 : 'FILE_DEVICE_DFS_FILE_SYSTEM',
    0x00000036 : 'FILE_DEVICE_DFS_VOLUME',
    0x00000007 : 'FILE_DEVICE_DISK',
    0x00000008 : 'FILE_DEVICE_DISK_FILE_SYSTEM',
    0x00000033 : 'FILE_DEVICE_DVD',
    0x00000009 : 'FILE_DEVICE_FILE_SYSTEM',
    0x0000003a : 'FILE_DEVICE_FIPS',
    0x00000034 : 'FILE_DEVICE_FULLSCREEN_VIDEO',
    0x0000000a : 'FILE_DEVICE_INPORT_PORT',
    0x0000000b : 'FILE_DEVICE_KEYBOARD',
    0x0000002f : 'FILE_DEVICE_KS',
    0x00000039 : 'FILE_DEVICE_KSEC',
    0x0000000c : 'FILE_DEVICE_MAILSLOT',
    0x0000002d : 'FILE_DEVICE_MASS_STORAGE',
    0x0000000d : 'FILE_DEVICE_MIDI_IN',
    0x0000000e : 'FILE_DEVICE_MIDI_OUT',
    0x0000002b : 'FILE_DEVICE_MODEM',
    0x0000000f : 'FILE_DEVICE_MOUSE',
    0x00000010 : 'FILE_DEVICE_MULTI_UNC_PROVIDER',
    0x00000011 : 'FILE_DEVICE_NAMED_PIPE',
    0x00000012 : 'FILE_DEVICE_NETWORK',
    0x00000013 : 'FILE_DEVICE_NETWORK_BROWSER',
    0x00000014 : 'FILE_DEVICE_NETWORK_FILE_SYSTEM',
    0x00000028 : 'FILE_DEVICE_NETWORK_REDIRECTOR',
    0x00000015 : 'FILE_DEVICE_NULL',
    0x00000016 : 'FILE_DEVICE_PARALLEL_PORT',
    0x00000017 : 'FILE_DEVICE_PHYSICAL_NETCARD',
    0x00000018 : 'FILE_DEVICE_PRINTER',
    0x00000019 : 'FILE_DEVICE_SCANNER',
    0x0000001c : 'FILE_DEVICE_SCREEN',
    0x00000037 : 'FILE_DEVICE_SERENUM',
    0x0000001a : 'FILE_DEVICE_SERIAL_MOUSE_PORT',
    0x0000001b : 'FILE_DEVICE_SERIAL_PORT',
    0x00000031 : 'FILE_DEVICE_SMARTCARD',
    0x0000002e : 'FILE_DEVICE_SMB',
    0x0000001d : 'FILE_DEVICE_SOUND',
    0x0000001e : 'FILE_DEVICE_STREAMS',
    0x0000001f : 'FILE_DEVICE_TAPE',
    0x00000020 : 'FILE_DEVICE_TAPE_FILE_SYSTEM',
    0x00000038 : 'FILE_DEVICE_TERMSRV',
    0x00000021 : 'FILE_DEVICE_TRANSPORT',
    0x00000022 : 'FILE_DEVICE_UNKNOWN',
    0x0000002c : 'FILE_DEVICE_VDM',
    0x00000023 : 'FILE_DEVICE_VIDEO',
    0x00000024 : 'FILE_DEVICE_VIRTUAL_DISK',
    0x00000025 : 'FILE_DEVICE_WAVE_IN',
    0x00000026 : 'FILE_DEVICE_WAVE_OUT',
}

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

class _DRIVER_OBJECT(obj.CType, windows.ExecutiveObjectMixin):
    "Class for driver objects"

    def devices(self):
        "Enumerate the driver's device objects"

        device = self.DeviceObject.dereference()
        while device:
            yield device
            device = device.NextDevice.dereference()

    def is_valid(self):
        return (obj.CType.is_valid(self) and 
                        self.DriverStart % 0x1000 == 0)

class _DEVICE_OBJECT(obj.CType, windows.ExecutiveObjectMixin):
    "Class for device objects"

    def attached_devices(self):
        "Enumerate the device's attachees"

        device = self.AttachedDevice.dereference()
        while device:
            yield device
            device = device.AttachedDevice.dereference()

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

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

#--------------------------------------------------------------------------------
# devicetree plugin
#--------------------------------------------------------------------------------

class DeviceTree(filescan.DriverScan):
    "Show device tree"

    def render_text(self, outfd, data):
        for driver in data:
            header = driver.get_object_header()

            outfd.write("DRV 0x{0:08x} {1}\n".format(driver.obj_offset,
                str(driver.DriverName or header.NameInfo.Name or '')))

            for device in driver.devices():

                device_header = obj.Object("_OBJECT_HEADER", offset = device.obj_offset -
                        device.obj_vm.profile.get_obj_offset("_OBJECT_HEADER", "Body"),
                        vm = device.obj_vm,
                        native_vm = device.obj_native_vm
                        )

                device_name = str(device_header.NameInfo.Name or '')

                outfd.write("---| DEV {0:#x} {1} {2}\n".format(
                    device.obj_offset,
                    device_name,
                    DEVICE_CODES.get(device.DeviceType.v(), "UNKNOWN")))

                level = 0

                for att_device in device.attached_devices():

                    device_header = obj.Object("_OBJECT_HEADER", offset = att_device.obj_offset -
                        att_device.obj_vm.profile.get_obj_offset("_OBJECT_HEADER", "Body"),
                        vm = att_device.obj_vm,
                        native_vm = att_device.obj_native_vm
                        )

                    device_name = str(device_header.NameInfo.Name or '')
                    name = (device_name + " - " +
                           str(att_device.DriverObject.DriverName or ''))

                    outfd.write("------{0}| ATT {1:#x} {2} {3}\n".format(
                        "---" * level,
                        att_device.obj_offset,
                        name,
                        DEVICE_CODES.get(att_device.DeviceType.v(), "UNKNOWN")))

                    level += 1

#--------------------------------------------------------------------------------
# driverirp plugin
#--------------------------------------------------------------------------------

class DriverIrp(filescan.DriverScan):
    "Driver IRP hook detection"

    def __init__(self, config, *args, **kwargs):
        filescan.DriverScan.__init__(self, config, *args, **kwargs)
        config.add_option("REGEX", short_option = 'r', type = 'str',
                          action = 'store',
                          help = 'Analyze drivers matching REGEX')

    def render_text(self, outfd, data):

        addr_space = utils.load_as(self._config)

        # Compile the regular expression for filtering by driver name 
        if self._config.regex != None:
            mod_re = re.compile(self._config.regex, re.I)
        else:
            mod_re = None

        mods = dict((addr_space.address_mask(mod.DllBase), mod) for mod in modules.lsmod(addr_space))
        mod_addrs = sorted(mods.keys())

        bits = addr_space.profile.metadata.get('memory_model', '32bit')

        self.table_header(None, [('i', ">4"),
                                 ('Funcs', "36"),
                                 ('addr', '[addrpad]'),
                                 ('name', '')
                                 ])

        for driver in data:

            header = driver.get_object_header()

            driver_name = str(header.NameInfo.Name or '')
            # Continue if a regex was supplied and it doesn't match 
            if mod_re != None:
                if not (mod_re.search(driver_name) or
                        mod_re.search(driver_name)): continue

            # Write the standard header for each driver object 
            outfd.write("{0}\n".format("-" * 50))
            outfd.write("DriverName: {0}\n".format(driver_name))
            outfd.write("DriverStart: {0:#x}\n".format(driver.DriverStart))
            outfd.write("DriverSize: {0:#x}\n".format(driver.DriverSize))
            outfd.write("DriverStartIo: {0:#x}\n".format(driver.DriverStartIo))

            # Write the address and owner of each IRP function 
            for i, function in enumerate(driver.MajorFunction):
                function = driver.MajorFunction[i]
                module = tasks.find_module(mods, mod_addrs, addr_space.address_mask(function))
                if module:
                    module_name = str(module.BaseDllName or '')
                else:
                    module_name = "Unknown"
                # This is where we check for inline hooks once the 
                # ApiHooks plugin is ported to 2.1. 
                self.table_row(outfd, i, MAJOR_FUNCTIONS[i], function, module_name)

                if self._config.verbose:
                    data = addr_space.zread(function, 64)
                    outfd.write("\n".join(
                        ["{0:#x} {1:<16} {2}".format(o, h, i)
                        for o, i, h in malfind.Disassemble(data = data, 
                            start = function, bits = bits, stoponret = True)
                    ]))
                    outfd.write("\n")

    def get_json(self, data):
        results = {}

        addr_space = utils.load_as(self._config)

        # Compile the regular expression for filtering by driver name
        if self._config.regex != None:
            mod_re = re.compile(self._config.regex, re.I)
        else:
            mod_re = None

        mods = dict((mod.DllBase, mod) for mod in modules.lsmod(addr_space))
        mod_addrs = sorted(mods.keys())

        bits = addr_space.profile.metadata.get('memory_model', '32bit')

        for driver in data:
            header = driver.get_object_header()

            driver_name = str(header.NameInfo.Name or '')

            # Continue if a regex was supplied and it doesn't match
            if mod_re != None:
                if not (mod_re.search(driver_name) or
                        mod_re.search(driver_name)): continue

            row = {'DriverName':driver_name,
                   'DriverStart':int(driver.DriverStart),
                   'DriverSize':int(driver.DriverSize),
                   'DriverStartIo':int(driver.DriverStartIo),
                   'IrpFunctions':[],
                   'Devices':[],
                   }

            # Write the address and owner of each IRP function
            for i, function in enumerate(driver.MajorFunction):
                function = driver.MajorFunction[i]
                module = tasks.find_module(mods, mod_addrs, function)
                irp = {'index':int(i),
                        'FunctionName':str(MAJOR_FUNCTIONS[i]),
                        'FunctionAddress':int(function),
                        'Disassembly':[]
                        }
                if module:
                    irp['BaseDllName'] = str(module.BaseDllName or '')
                else:
                    irp['BaseDllName'] = ''


                if self._config.verbose:
                    data = addr_space.zread(function, 64)
                    irp["Disassembly"] = [{'Address':int(o), "Bytes":str(h), "Instruction":str(i)}
                        for o, i, h in malfind.Disassemble(data = data,
                            start = function, bits = bits, stoponret = True)
                    ]
                row['IrpFunctions'].append(irp)

            for device in driver.devices():

                device_header = obj.Object("_OBJECT_HEADER", offset = device.obj_offset -
                        device.obj_vm.profile.get_obj_offset("_OBJECT_HEADER", "Body"),
                        vm = device.obj_vm,
                        native_vm = device.obj_native_vm
                        )

                device_name = str(device_header.NameInfo.Name or '')

                dev = {'Offset': int(device.obj_offset),
                          'DeviceName': device_name,
                          'DeviceCodes': DEVICE_CODES.get(device.DeviceType.v(), "UNKNOWN"),
                          'AttachedDevices': []}

                level = 0

                for att_device in device.attached_devices():

                    device_header = obj.Object("_OBJECT_HEADER", offset = att_device.obj_offset -
                        att_device.obj_vm.profile.get_obj_offset("_OBJECT_HEADER", "Body"),
                        vm = att_device.obj_vm,
                        native_vm = att_device.obj_native_vm
                        )

                    device_name = str(device_header.NameInfo.Name or '')
                    name = (device_name + " - " +
                           str(att_device.DriverObject.DriverName or ''))

                    attached_device = {'Offset': att_device.obj_offset,
                                       'DeviceName': name,
                                       'DeviceCodes': DEVICE_CODES.get(att_device.DeviceType.v(), "UNKNOWN"),
                                       'Level': level}

                    level += 1
                    dev['AttachedDevices'].append(attached_device)
                row['Devices'].append(dev)

            _addr = str(row['DriverStart'])
            results [_addr ] = row

        return results

    def render_json(self, outfd, data):
        outfd.write( json.dumps(self.get_json(data), indent=4) + '\n')