# This file is Copyright 2020 Volatility Foundation and licensed under the Volatility Software License 1.0
# which is available at https://www.volatilityfoundation.org/license/vsl-v1.0
#
import contextlib
import logging
from typing import List

from volatility3.framework import exceptions, renderers, interfaces
from volatility3.framework.configuration import requirements
from volatility3.framework.renderers import format_hints
from volatility3.plugins.windows import pslist

vollog = logging.getLogger(__name__)


class Memmap(interfaces.plugins.PluginInterface):
    """Prints the memory map"""

    _required_framework_version = (2, 0, 0)

    @classmethod
    def get_requirements(cls) -> List[interfaces.configuration.RequirementInterface]:
        # Since we're calling the plugin, make sure we have the plugin's requirements
        return [
            requirements.ModuleRequirement(name = 'kernel', description = 'Windows kernel',
                                                     architectures = ["Intel32", "Intel64"]),
            requirements.PluginRequirement(name = 'pslist', plugin = pslist.PsList, version = (2, 0, 0)),
            requirements.IntRequirement(name = 'pid',
                                        description = "Process ID to include (all other processes are excluded)",
                                        optional = True),
            requirements.BooleanRequirement(name = 'dump',
                                            description = "Extract listed memory segments",
                                            default = False,
                                            optional = True)
        ]


    def _generator(self, procs):
        for proc in procs:
            pid = "Unknown"

            try:
                pid = proc.UniqueProcessId
                proc_layer_name = proc.add_process_layer()
                proc_layer = self.context.layers[proc_layer_name]
            except exceptions.InvalidAddressException as excp:
                vollog.debug("Process {}: invalid address {} in layer {}".format(pid, excp.invalid_address,
                                                                                 excp.layer_name))
                continue

            if self.config['dump']:
                file_handle = self.open(f"pid.{pid}.dmp")
            else:
                # Ensure the file isn't actually created if not needed
                file_handle = contextlib.ExitStack()
            with file_handle as file_data:
                file_offset = 0
                for mapval in proc_layer.mapping(0x0, proc_layer.maximum_address, ignore_errors = True):
                    offset, size, mapped_offset, mapped_size, maplayer = mapval

                    file_output = "Disabled"
                    if self.config['dump']:
                        try:
                            data = proc_layer.read(offset, size, pad = True)
                            file_data.write(data)
                            file_output = file_handle.preferred_filename
                        except exceptions.InvalidAddressException:
                            file_output = "Error outputting to file"
                            vollog.debug("Unable to write {}'s address {} to {}".format(
                                proc_layer_name, offset, file_handle.preferred_filename))

                    yield (0, (format_hints.Hex(offset), format_hints.Hex(mapped_offset),
                               format_hints.Hex(mapped_size),
                               format_hints.Hex(file_offset), file_output))

                    file_offset += mapped_size
                    offset += mapped_size

    def run(self):
        filter_func = pslist.PsList.create_pid_filter([self.config.get('pid', None)])
        kernel = self.context.modules[self.config['kernel']]

        return renderers.TreeGrid([("Virtual", format_hints.Hex), ("Physical", format_hints.Hex),
                                   ("Size", format_hints.Hex), ("Offset in File", format_hints.Hex),
                                   ("File output", str)],
                                  self._generator(
                                      pslist.PsList.list_processes(context = self.context,
                                                                   layer_name = kernel.layer_name,
                                                                   symbol_table = kernel.symbol_table_name,
                                                                   filter_func = filter_func)))
