"""A parser for Portable Executable (PE) files.

Also see:
  https://cloud.google.com/blog/topics/threat-intelligence/tracking-malware-import-hashing/
"""

import os

import pefile

from dfdatetime import posix_time as dfdatetime_posix_time
from dfvfs.helpers import data_slice as dfvfs_data_slice
from dtfabric.runtime import data_maps as dtfabric_data_maps

from plaso.containers import artifacts
from plaso.containers import events
from plaso.helpers.windows import languages
from plaso.helpers.windows import resource_files
from plaso.lib import errors
from plaso.lib import dtfabric_helper
from plaso.lib import specification
from plaso.parsers import interface
from plaso.parsers import manager


class PEDLLImportEventData(events.EventData):
    """Portable Executable (PE) DLL import event data.

    Attributes:
      delayed_import (bool): True if the DLL is imported at run-time.
      modification_time (dfdatetime.DateTimeValues): last modification date and time.
      name (str): name of the imported DLL.
    """

    DATA_TYPE = "pe_coff:dll_import"

    def __init__(self):
        """Initializes event data."""
        super().__init__(data_type=self.DATA_TYPE)
        self.delayed_import = None
        self.modification_time = None
        self.name = None


class PEFileEventData(events.EventData):
    """Portable Executable (PE) file event data.

    Attributes:
      creation_time (dfdatetime.DateTimeValues): creation date and time.
      export_dll_name (str): name of the exported DLL.
      export_table_modification_time (dfdatetime.DateTimeValues): export table last
          modification date and time.
      imphash (str): "Import Hash" of the Portable Executable (PE) file.
      load_configuration_table_modification_time (dfdatetime.DateTimeValues): load
          configuration table last modification date and time.
      pe_type (str): type of Portable Executable (PE) file.
      section_names (list[str]): names of the sections in the Portable Executable (PE)
          file.
    """

    DATA_TYPE = "pe_coff:file"

    def __init__(self):
        """Initializes event data."""
        super().__init__(data_type=self.DATA_TYPE)
        self.creation_time = None
        self.export_dll_name = None
        self.export_table_modification_time = None
        self.imphash = None
        self.load_configuration_table_modification_time = None
        self.pe_type = None
        self.section_names = None


class PEResourceEventData(events.EventData):
    """Portable Executable (PE) resource event data.

    Attributes:
      identifier (int): identifier of the resource.
      modification_time (dfdatetime.DateTimeValues): last modification date and time.
      name (str): name of the resource.
    """

    DATA_TYPE = "pe_coff:resource"

    def __init__(self):
        """Initializes event data."""
        super().__init__(data_type=self.DATA_TYPE)
        self.identifier = None
        self.modification_time = None
        self.name = None


class PEParser(interface.FileObjectParser, dtfabric_helper.DtFabricHelper):
    """Parser for Portable Executable (PE) files."""

    NAME = "pe"
    DATA_FORMAT = "Portable Executable (PE) file"

    _DEFINITION_FILE = os.path.join(os.path.dirname(__file__), "pe_resources.yaml")

    _PE_DIRECTORIES = [
        pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT"],
        pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_EXPORT"],
        pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_IMPORT"],
        pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG"],
        pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_RESOURCE"],
    ]

    def __init__(self):
        """Initializes a PE parser."""
        super().__init__()
        self._resource_file_helper = resource_files.WindowsResourceFileHelper

    def _GetPEType(self, pefile_object):
        """Retrieves the type of the PE file.

        Args:
          pefile_object (pefile.PE): pefile object.

        Returns:
          str: type of the Portable Executable (PE) file.
        """
        if pefile_object.is_dll():
            return "Dynamic Link Library (DLL)"
        if pefile_object.is_exe():
            return "Executable (EXE)"
        if pefile_object.is_driver():
            return "Driver (SYS)"
        return "Unknown PE type"

    def _GetSectionNames(self, pefile_object):
        """Retrieves all PE section names.

        Args:
          pefile_object (pefile.PE): pefile object.

        Returns:
          list[str]: names of the sections.
        """
        section_names = []
        for section in pefile_object.sections:
            section_name = getattr(section, "Name", b"")
            # Ensure the name is decoded correctly.
            try:
                section_name = section_name.decode("unicode_escape")
            except UnicodeDecodeError:
                section_name = repr(section_name)
            section_names.append(section_name)

        return section_names

    def _ParseLoadConfigurationTable(self, pefile_object, event_data):
        """Parses the load configuration table.

        Args:
          pefile_object (pefile.PE): pefile object.
          event_data (PEFileEventData): event data.
        """
        load_configuration_table = getattr(
            pefile_object, "DIRECTORY_ENTRY_LOAD_CONFIG", None
        )
        if not load_configuration_table:
            return

        timestamp = getattr(load_configuration_table.struct, "TimeDateStamp", None)
        if timestamp:
            date_time = dfdatetime_posix_time.PosixTime(timestamp=timestamp)
            event_data.load_configuration_table_modification_time = date_time

    def _ParseDelayImportTable(self, parser_mediator, pefile_object):
        """Parses the delay import table.

        Args:
          parser_mediator (ParserMediator): mediates interactions between parsers and
              other components, such as storage and dfVFS.
          pefile_object (pefile.PE): pefile object.
        """
        delay_import_table = getattr(
            pefile_object, "DIRECTORY_ENTRY_DELAY_IMPORT", None
        )
        if not delay_import_table:
            return

        for table_entry in delay_import_table:
            timestamp = getattr(table_entry.struct, "dwTimeStamp", None)
            if not timestamp:
                continue

            corrupted = False

            dll_name = getattr(table_entry, "dll", "")
            try:
                dll_name = dll_name.decode("ascii")
            except UnicodeDecodeError:
                parser_mediator.ProduceWarning(
                    "unable to decode delay import table DLL name as ASCII. "
                    "Unsupported code points are escaped."
                )
                dll_name = dll_name.decode("ascii", errors="backslashreplace")
                corrupted = True

            event_data = PEDLLImportEventData()
            event_data.delayed_import = True
            event_data.name = dll_name or None

            date_time = dfdatetime_posix_time.PosixTime(timestamp=timestamp)
            event_data.modification_time = date_time

            parser_mediator.ProduceEventData(event_data, corrupted=corrupted)

    def _ParseExportTable(self, parser_mediator, pefile_object, event_data):
        """Parses the export table.

        Args:
          parser_mediator (ParserMediator): mediates interactions between parsers and
              other components, such as storage and dfVFS.
          pefile_object (pefile.PE): pefile object.
          event_data (PEFileEventData): event data.

        Returns:
          bool: value to indicate the export table was corrupted.
        """
        export_table = getattr(pefile_object, "DIRECTORY_ENTRY_EXPORT", None)
        if not export_table:
            return False

        corrupted = False

        dll_name = getattr(export_table, "name", "")
        try:
            dll_name = dll_name.decode("ascii")
        except UnicodeDecodeError:
            parser_mediator.ProduceWarning(
                "unable to decode export table DLL name as ASCII. Unsupported code "
                "points are escaped."
            )
            dll_name = dll_name.decode("ascii", errors="backslashreplace")
            corrupted = True

        event_data.export_dll_name = dll_name or None

        timestamp = getattr(export_table.struct, "TimeDateStamp", None)
        if timestamp:
            date_time = dfdatetime_posix_time.PosixTime(timestamp=timestamp)
            event_data.export_table_modification_time = date_time

        return corrupted

    def _ParseImportTable(self, parser_mediator, pefile_object):
        """Parses the import table.

        Args:
          parser_mediator (ParserMediator): mediates interactions between parsers and
              other components, such as storage and dfVFS.
          pefile_object (pefile.PE): pefile object.
        """
        import_table = getattr(pefile_object, "DIRECTORY_ENTRY_IMPORT", None)
        if not import_table:
            return

        for table_entry in import_table:
            timestamp = getattr(table_entry.struct, "TimeDateStamp", None)
            if not timestamp:
                continue

            corrupted = False

            dll_name = getattr(table_entry, "dll", "")
            try:
                dll_name = dll_name.decode("ascii")
            except UnicodeDecodeError:
                parser_mediator.ProduceWarning(
                    "unable to decode import table DLL name as ASCII. Unsupported code "
                    "points are escaped."
                )
                dll_name = dll_name.decode("ascii", errors="backslashreplace")
                corrupted = True

            event_data = PEDLLImportEventData()
            event_data.delayed_import = False
            event_data.name = dll_name or None

            date_time = dfdatetime_posix_time.PosixTime(timestamp=timestamp)
            event_data.modification_time = date_time

            parser_mediator.ProduceEventData(event_data, corrupted=corrupted)

    def _ParseResourceSection(self, parser_mediator, pefile_object):
        """Parses the resource section.

        Args:
          parser_mediator (ParserMediator): mediates interactions between parsers and
              other components, such as storage and dfVFS.
          pefile_object (pefile.PE): pefile object.
        """
        resources = getattr(pefile_object, "DIRECTORY_ENTRY_RESOURCE", None)
        if not resources:
            return

        message_table_resource = None
        winevt_template_resource = None
        for resource in resources.entries:
            if resource.name:
                resource_name = str(resource.name)
            else:
                resource_name = None

            timestamp = getattr(resource.directory, "TimeDateStamp", None)
            if timestamp:
                event_data = PEResourceEventData()
                event_data.identifier = resource.id
                event_data.name = resource_name

                date_time = dfdatetime_posix_time.PosixTime(timestamp=timestamp)
                event_data.modification_time = date_time

                parser_mediator.ProduceEventData(event_data)

            # Only extract message strings from the first message table resource.
            if resource.id == 11 and not message_table_resource:
                message_table_resource = resource

            # Need to cast resource.name name to a string since it is of type
            # pefile.UnicodeStringWrapperPostProcessor and would fail the comparison
            # with 'WEVT_TEMPLATE' otherwise.
            elif resource_name == "WEVT_TEMPLATE" and not winevt_template_resource:
                winevt_template_resource = resource

        if parser_mediator.extract_winevt_resources:
            # Only extract message strings from EventLog message files.
            message_file = parser_mediator.GetWindowsEventLogMessageFile()
            if message_file:
                parser_mediator.AddWindowsEventLogMessageFile(message_file)

                if message_table_resource:
                    self._ParseMessageTableResource(
                        parser_mediator,
                        pefile_object,
                        message_file,
                        message_table_resource,
                    )

                if winevt_template_resource:
                    self._ParseWevtTemplateResource(
                        parser_mediator,
                        pefile_object,
                        message_file,
                        winevt_template_resource,
                    )

    def _ParseMessageTable(self, parser_mediator, message_table, data):
        """Parses a message table.

        Args:
          parser_mediator (ParserMediator): mediates interactions between parsers and
              other components, such as storage and dfVFS.
          message_table (WindowsEventLogMessageTableArtifact): Windows EventLog message
              table.
          data (bytes): message table data.

        Raises:
          ParseError: when the message table cannot be parsed.
        """
        message_table_identifier = message_table.GetIdentifier()

        message_table_header_map = self._GetDataTypeMap("message_table_header")
        message_table_entry_map = self._GetDataTypeMap("message_table_entry")
        message_table_string_map = self._GetDataTypeMap("message_table_string")

        data_offset = 0

        context = dtfabric_data_maps.DataTypeMapContext()

        try:
            message_table_header = self._ReadStructureFromByteStream(
                data, data_offset, message_table_header_map, context=context
            )
        except (ValueError, errors.ParseError) as exception:
            raise errors.ParseError(
                f"Unable to read message table header with error: {exception!s}"
            )

        data_offset += context.byte_size

        for entry_index in range(message_table_header.number_of_entries):
            context = dtfabric_data_maps.DataTypeMapContext()

            try:
                message_table_entry = self._ReadStructureFromByteStream(
                    data[data_offset:],
                    data_offset,
                    message_table_entry_map,
                    context=context,
                )
            except (ValueError, errors.ParseError) as exception:
                raise errors.ParseError(
                    f"Unable to read message table entry: {entry_index:d} at offset: "
                    f"{data_offset:d} with error: {exception!s}"
                )

            data_offset += context.byte_size

            message_identifier = message_table_entry.first_message_identifier
            string_offset = message_table_entry.first_string_offset
            while message_identifier <= message_table_entry.last_message_identifier:
                try:
                    message_table_string = self._ReadStructureFromByteStream(
                        data[string_offset:], string_offset, message_table_string_map
                    )
                except (ValueError, errors.ParseError) as exception:
                    raise errors.ParseError(
                        f"Unable to read message table string: "
                        f"0x{message_identifier:08x} at offset: {string_offset:d} with "
                        f"error: {exception!s}"
                    )

                if message_table_string.flags & 0x01:
                    string_encoding = "utf-16-le"
                else:
                    string_encoding = parser_mediator.GetCodePage()

                try:
                    string = message_table_string.data.decode(string_encoding)
                except UnicodeDecodeError:
                    raise errors.ParseError(
                        f"Unable to decode {string_encoding:s} encoded message table "
                        f"string: 0x{message_identifier:08x} at offset: "
                        f"{string_offset:d} with error: {exception!s}"
                    )

                string_offset += message_table_string.data_size

                _, alignment_padding = divmod(string_offset, 4)
                if alignment_padding > 0:
                    string_offset += 4 - alignment_padding

                string = self._resource_file_helper.FormatMessageStringInPEP3101(string)

                message_string = artifacts.WindowsEventLogMessageStringArtifact(
                    message_identifier=message_identifier,
                    string=string,
                )
                message_string.SetMessageTableIdentifier(message_table_identifier)
                parser_mediator.AddWindowsEventLogMessageString(message_string)

                message_identifier += 1

    def _ParseMessageTableResource(
        self, parser_mediator, pefile_object, message_file, message_table_resource
    ):
        """Parses a message table resource.

        Args:
          parser_mediator (ParserMediator): mediates interactions between parsers
              and other components, such as storage and dfVFS.
          pefile_object (pefile.PE): pefile object.
          message_file (WindowsEventLogMessageFileArtifact): Windows EventLog message
              file.
          message_table_resource (pefile.ResourceDirEntryData): message table resource.
        """
        if (
            not message_table_resource
            or not message_table_resource.directory
            or not message_table_resource.directory.entries
            or not message_table_resource.directory.entries[0].directory
        ):
            return

        message_file_identifier = message_file.GetIdentifier()

        desired_language_tag = parser_mediator.GetLanguageTag().lower()

        for entry in message_table_resource.directory.entries[0].directory.entries:
            language_tag = languages.WindowsLanguageHelper.GetLanguageTagForLCID(
                entry.id
            )
            # TODO: add support for common language tag fallback.
            if not language_tag or language_tag.lower() != desired_language_tag:
                continue

            message_table = artifacts.WindowsEventLogMessageTableArtifact(
                language_identifier=entry.id,
            )
            message_table.SetMessageFileIdentifier(message_file_identifier)
            parser_mediator.AddWindowsEventLogMessageTable(message_table)

            # TODO: use file offset?
            offset = getattr(entry.data.struct, "OffsetToData", None)
            size = getattr(entry.data.struct, "Size", None)
            data = pefile_object.get_memory_mapped_image()[offset : offset + size]

            self._ParseMessageTable(parser_mediator, message_table, data)

    def _ParseWevtTemplate(self, parser_mediator, message_file, data):
        """Parses a WEVT_TEMPLATE.

        Args:
          parser_mediator (ParserMediator): mediates interactions between parsers
              and other components, such as storage and dfVFS.
          message_file (WindowsEventLogMessageFileArtifact): Windows EventLog message
              file.
          data (bytes): message table data.

        Raises:
          ParseError: when the message table cannot be parsed.
        """
        message_file_identifier = message_file.GetIdentifier()

        wevt_manifest_map = self._GetDataTypeMap("wevt_instrumentation_manifest")
        wevt_event_provider_map = self._GetDataTypeMap("wevt_event_provider")
        wevt_event_definitions_map = self._GetDataTypeMap("wevt_event_definitions")

        try:
            manifest = self._ReadStructureFromByteStream(data, 0, wevt_manifest_map)
        except (ValueError, errors.ParseError) as exception:
            raise errors.ParseError(
                f"Unable to read WEVT instrumentation manifest with error: "
                f"{exception!s}"
            )

        for event_provider_descriptor in manifest.event_provider_descriptors:
            data_offset = event_provider_descriptor.data_offset
            try:
                event_provider = self._ReadStructureFromByteStream(
                    data[data_offset:], data_offset, wevt_event_provider_map
                )
            except (ValueError, errors.ParseError) as exception:
                raise errors.ParseError(
                    f"Unable to read WEVT event provider with error: {exception!s}"
                )

            provider_identifier = (
                f"{{{event_provider_descriptor.provider_identifier!s}}}"
            )
            for (
                provider_element_descriptor
            ) in event_provider.provider_element_descriptors:
                data_offset = provider_element_descriptor.data_offset

                # Only extract event definitions.
                if data[data_offset : data_offset + 4] == b"EVNT":
                    try:
                        event_definitions = self._ReadStructureFromByteStream(
                            data[data_offset:], data_offset, wevt_event_definitions_map
                        )
                    except (ValueError, errors.ParseError) as exception:
                        raise errors.ParseError(
                            f"Unable to read WEVT event definitions with error: "
                            f"{exception!s}"
                        )

                    for event_definition in event_definitions.definitions:
                        event_definition = artifacts.WindowsWevtTemplateEvent(
                            identifier=event_definition.identifier,
                            message_identifier=event_definition.message_identifier,
                            provider_identifier=provider_identifier,
                            version=event_definition.version,
                        )
                        event_definition.SetMessageFileIdentifier(
                            message_file_identifier
                        )
                        parser_mediator.AddWindowsWevtTemplateEvent(event_definition)

    def _ParseWevtTemplateResource(
        self, parser_mediator, pefile_object, message_file, wevt_template_resource
    ):
        """Parses a WEVT_TEMPLATE resource.

        Args:
          parser_mediator (ParserMediator): mediates interactions between parsers
              and other components, such as storage and dfVFS.
          pefile_object (pefile.PE): pefile object.
          message_file (WindowsEventLogMessageFileArtifact): Windows EventLog message
              file.
          wevt_template_resource (pefile.ResourceDirEntryData): WEVT_TEMPLATE resource.

        Raises:
          ParseError: when the message table cannot be parsed.
        """
        if (
            not wevt_template_resource
            or not wevt_template_resource.directory
            or not wevt_template_resource.directory.entries
            or not wevt_template_resource.directory.entries[0].directory
        ):
            return

        desired_language_tag = parser_mediator.GetLanguageTag().lower()

        for entry in wevt_template_resource.directory.entries[0].directory.entries:
            language_tag = languages.WindowsLanguageHelper.GetLanguageTagForLCID(
                entry.id
            )
            # TODO: add support for common language tag fallback.
            if not language_tag or language_tag.lower() != desired_language_tag:
                continue

            # TODO: use file offset?
            offset = getattr(entry.data.struct, "OffsetToData", None)
            size = getattr(entry.data.struct, "Size", None)
            data = pefile_object.get_memory_mapped_image()[offset : offset + size]

            self._ParseWevtTemplate(parser_mediator, message_file, data)

    @classmethod
    def GetFormatSpecification(cls):
        """Retrieves the format specification."""
        format_specification = specification.FormatSpecification(cls.NAME)
        format_specification.AddNewSignature(b"MZ", offset=0)
        return format_specification

    def ParseFileObject(self, parser_mediator, file_object):
        """Parses a Portable Executable (PE) file-like object.

        Args:
          parser_mediator (ParserMediator): mediates interactions between parsers
              and other components, such as storage and dfVFS.
          file_object (dfvfs.FileIO): a file-like object.

        Raises:
          WrongParser: when the file cannot be parsed.
        """
        pe_data_slice = dfvfs_data_slice.DataSlice(file_object)
        try:
            pefile_object = pefile.PE(data=pe_data_slice, fast_load=True)
            pefile_object.parse_data_directories(directories=self._PE_DIRECTORIES)
        except Exception as exception:
            raise errors.WrongParser(
                f"Unable to read PE file with error: {exception!s}"
            )

        event_data = PEFileEventData()
        # Note that the result of get_imphash() is an empty string if there is no
        # import hash.
        event_data.imphash = pefile_object.get_imphash() or None
        event_data.pe_type = self._GetPEType(pefile_object)
        event_data.section_names = self._GetSectionNames(pefile_object)

        timestamp = getattr(pefile_object.FILE_HEADER, "TimeDateStamp", None)
        if timestamp:
            event_data.creation_time = dfdatetime_posix_time.PosixTime(
                timestamp=timestamp
            )

        corrupted = self._ParseExportTable(parser_mediator, pefile_object, event_data)

        self._ParseLoadConfigurationTable(pefile_object, event_data)

        parser_mediator.ProduceEventData(event_data, corrupted=corrupted)

        self._ParseImportTable(parser_mediator, pefile_object)

        self._ParseDelayImportTable(parser_mediator, pefile_object)

        self._ParseResourceSection(parser_mediator, pefile_object)


manager.ParsersManager.RegisterParser(PEParser)
