<commit_before>import os
import struct
import logging
import archinfo
import pefile
from .symbol import WinSymbol
from .regions import PESection
from .relocation.generic import DllImport, IMAGE_REL_BASED_HIGHADJ, IMAGE_REL_BASED_ABSOLUTE
from .relocation import get_relocation
from .. import register_backend, Backend
from ...address_translator import AT


l = logging.getLogger('cle.pe')


class PE(Backend):
    """
    Representation of a PE (i.e. Windows) binary.
    """

    def __init__(self, *args, **kwargs):
        super(PE, self).__init__(*args, **kwargs)
        self.segments = self.sections # in a PE, sections and segments have the same meaning
        self.os = 'windows'
        if self.binary is None:
            self._pe = pefile.PE(data=self.binary_stream.read())
        elif self.binary in self._pefile_cache: # these objects are not mutated, so they are reusable within a process
            self._pe = self._pefile_cache[self.binary]
        else:
            self._pe = pefile.PE(self.binary)
            if not self.is_main_bin:
                # only cache shared libraries, the main binary will not be reused
                self._pefile_cache[self.binary] = self._pe

        if self.arch is None:
            self.set_arch(archinfo.arch_from_id(pefile.MACHINE_TYPE[self._pe.FILE_HEADER.Machine]))

        self.mapped_base = self.linked_base = self._pe.OPTIONAL_HEADER.ImageBase

        self._entry = AT.from_rva(self._pe.OPTIONAL_HEADER.AddressOfEntryPoint, self).to_lva()

        if hasattr(self._pe, 'DIRECTORY_ENTRY_IMPORT'):
            self.deps = [entry.dll.lower() for entry in self._pe.DIRECTORY_ENTRY_IMPORT]
        else:
            self.deps = []

        if self.binary is not None and not self.is_main_bin:
            self.provides = os.path.basename(self.binary).lower()
        else:
            self.provides = None

        self.tls_used = False
        self.tls_data_start = None
        self.tls_data_size = None
        self.tls_index_address = None
        self.tls_callbacks = None
        self.tls_size_of_zero_fill = None
        self.tls_module_id = None
        self.tls_data_pointer = None

        self._exports = {}
        self._ordinal_exports = {}
        self._symbol_cache = self._exports # same thing
        self._handle_imports()
        self._handle_exports()
        self.__register_relocs()
        self._register_tls()
        self._register_sections()

        self.linking = 'dynamic' if self.deps else 'static'
        self.supports_nx = self._pe.OPTIONAL_HEADER.DllCharacteristics & 0x100 != 0
        self.pic = self._pe.OPTIONAL_HEADER.DllCharacteristics & 0x40 != 0

        self.jmprel = self._get_jmprel()

        self.memory.add_backer(0, self._pe.get_memory_mapped_image())

    _pefile_cache = {}

    @staticmethod
    def is_compatible(stream):
        identstring = stream.read(0x1000)
        stream.seek(0)
        if identstring.startswith('MZ') and len(identstring) > 0x40:
            peptr = struct.unpack('I', identstring[0x3c:0x40])[0]
            if peptr < len(identstring) and identstring[peptr:peptr + 4] == 'PE\0\0':
                return True
        return False

    @classmethod
    def check_compatibility(cls, spec, obj):
        if hasattr(spec, 'read') and hasattr(spec, 'seek'):
            pe = pefile.PE(data=spec.read(), fast_load=True)
        else:
            pe = pefile.PE(spec, fast_load=True)

        arch = archinfo.arch_from_id(pefile.MACHINE_TYPE[pe.FILE_HEADER.Machine])
        return arch == obj.arch

    #
    # Public methods
    #

    def get_symbol(self, name):
        if name.startswith('ordinal.'):
            return self._ordinal_exports.get(int(name.split('.')[1]), None)
        return self._exports.get(name, None)

    #
    # Private methods
    #

    def _get_jmprel(self):
        return self.imports

    def _handle_imports(self):
        if hasattr(self._pe, 'DIRECTORY_ENTRY_IMPORT'):
            for entry in self._pe.DIRECTORY_ENTRY_IMPORT:
                for imp in entry.imports:
                    imp_name = imp.name
                    if imp_name is None: # must be an import by ordinal
                        imp_name = "ordinal.%d.%s" % (imp.ordinal, entry.dll.lower())

                    symb = WinSymbol(owner=self, name=imp_name, addr=0, is_import=True, is_export=False, ordinal_number=imp.ordinal, forwarder=None)
                    reloc = self._make_reloc(addr=AT.from_lva(imp.address, self).to_rva(), reloc_type=None, symbol=symb, resolvewith=entry.dll)

                    if reloc is not None:
                        self.imports[imp_name] = reloc
                        self.relocs.append(reloc)


    def _handle_exports(self):
        if hasattr(self._pe, 'DIRECTORY_ENTRY_EXPORT'):
            symbols = self._pe.DIRECTORY_ENTRY_EXPORT.symbols
            for exp in symbols:
                symb = WinSymbol(self, exp.name, exp.address, False, True, exp.ordinal, exp.forwarder)
                self._exports[exp.name] = symb
                self._ordinal_exports[exp.ordinal] = symb


                if exp.forwarder is not None:
                    forwardlib = exp.forwarder.split('.', 1)[0].lower() + '.dll'
                    if forwardlib not in self.deps:
                        self.deps.append(forwardlib)


    def __register_relocs(self):
        if not hasattr(self._pe, 'DIRECTORY_ENTRY_BASERELOC'):
            if self.pic:
                l.info('No relocations found in PIC binary')
            return

        for base_reloc in self._pe.DIRECTORY_ENTRY_BASERELOC:
            entry_idx = 0
            while entry_idx < len(base_reloc.entries):
                reloc_data = base_reloc.entries[entry_idx]
                if reloc_data.type == pefile.RELOCATION_TYPE['IMAGE_REL_BASED_HIGHADJ']: # special case, occupies 2 entries
                    if entry_idx == len(base_reloc.entries):
                        l.warning('PE contains corrupt base relocation table')
                        break

                    next_entry = base_reloc.entries[entry_idx]
                    entry_idx += 1
                    reloc = self._make_reloc(addr=reloc_data.rva, reloc_type=reloc_data.type, next_rva=next_entry.rva)
                else:
                    reloc = self._make_reloc(addr=reloc_data.rva, reloc_type=reloc_data.type)

                if reloc is not None:
                    self.relocs.append(reloc)

                entry_idx += 1

        return self.relocs

    def _make_reloc(self, addr, reloc_type, symbol=None, next_rva=None, resolvewith=None):

        # Handle special cases first<commit_msg>Change pic detection to work if there are relocations but DYNAMICBASE is not set<commit_after>import os
import struct
import logging
import archinfo
import pefile
from .symbol import WinSymbol
from .regions import PESection
from .relocation.generic import DllImport, IMAGE_REL_BASED_HIGHADJ, IMAGE_REL_BASED_ABSOLUTE
from .relocation import get_relocation
from .. import register_backend, Backend
from ...address_translator import AT


l = logging.getLogger('cle.pe')


class PE(Backend):
    """
    Representation of a PE (i.e. Windows) binary.
    """

    def __init__(self, *args, **kwargs):
        super(PE, self).__init__(*args, **kwargs)
        self.segments = self.sections # in a PE, sections and segments have the same meaning
        self.os = 'windows'
        if self.binary is None:
            self._pe = pefile.PE(data=self.binary_stream.read())
        elif self.binary in self._pefile_cache: # these objects are not mutated, so they are reusable within a process
            self._pe = self._pefile_cache[self.binary]
        else:
            self._pe = pefile.PE(self.binary)
            if not self.is_main_bin:
                # only cache shared libraries, the main binary will not be reused
                self._pefile_cache[self.binary] = self._pe

        if self.arch is None:
            self.set_arch(archinfo.arch_from_id(pefile.MACHINE_TYPE[self._pe.FILE_HEADER.Machine]))

        self.mapped_base = self.linked_base = self._pe.OPTIONAL_HEADER.ImageBase

        self._entry = AT.from_rva(self._pe.OPTIONAL_HEADER.AddressOfEntryPoint, self).to_lva()

        if hasattr(self._pe, 'DIRECTORY_ENTRY_IMPORT'):
            self.deps = [entry.dll.lower() for entry in self._pe.DIRECTORY_ENTRY_IMPORT]
        else:
            self.deps = []

        if self.binary is not None and not self.is_main_bin:
            self.provides = os.path.basename(self.binary).lower()
        else:
            self.provides = None

        self.tls_used = False
        self.tls_data_start = None
        self.tls_data_size = None
        self.tls_index_address = None
        self.tls_callbacks = None
        self.tls_size_of_zero_fill = None
        self.tls_module_id = None
        self.tls_data_pointer = None

        self.supports_nx = self._pe.OPTIONAL_HEADER.DllCharacteristics & 0x100 != 0
        self.pic = self._pe.OPTIONAL_HEADER.DllCharacteristics & 0x40 != 0

        self._exports = {}
        self._ordinal_exports = {}
        self._symbol_cache = self._exports # same thing
        self._handle_imports()
        self._handle_exports()
        self.__register_relocs()
        self._register_tls()
        self._register_sections()

        self.linking = 'dynamic' if self.deps else 'static'
        self.jmprel = self._get_jmprel()
        self.memory.add_backer(0, self._pe.get_memory_mapped_image())

    _pefile_cache = {}

    @staticmethod
    def is_compatible(stream):
        identstring = stream.read(0x1000)
        stream.seek(0)
        if identstring.startswith('MZ') and len(identstring) > 0x40:
            peptr = struct.unpack('I', identstring[0x3c:0x40])[0]
            if peptr < len(identstring) and identstring[peptr:peptr + 4] == 'PE\0\0':
                return True
        return False

    @classmethod
    def check_compatibility(cls, spec, obj):
        if hasattr(spec, 'read') and hasattr(spec, 'seek'):
            pe = pefile.PE(data=spec.read(), fast_load=True)
        else:
            pe = pefile.PE(spec, fast_load=True)

        arch = archinfo.arch_from_id(pefile.MACHINE_TYPE[pe.FILE_HEADER.Machine])
        return arch == obj.arch

    #
    # Public methods
    #

    def get_symbol(self, name):
        if name.startswith('ordinal.'):
            return self._ordinal_exports.get(int(name.split('.')[1]), None)
        return self._exports.get(name, None)

    #
    # Private methods
    #

    def _get_jmprel(self):
        return self.imports

    def _handle_imports(self):
        if hasattr(self._pe, 'DIRECTORY_ENTRY_IMPORT'):
            for entry in self._pe.DIRECTORY_ENTRY_IMPORT:
                for imp in entry.imports:
                    imp_name = imp.name
                    if imp_name is None: # must be an import by ordinal
                        imp_name = "ordinal.%d.%s" % (imp.ordinal, entry.dll.lower())

                    symb = WinSymbol(owner=self, name=imp_name, addr=0, is_import=True, is_export=False, ordinal_number=imp.ordinal, forwarder=None)
                    reloc = self._make_reloc(addr=AT.from_lva(imp.address, self).to_rva(), reloc_type=None, symbol=symb, resolvewith=entry.dll)

                    if reloc is not None:
                        self.imports[imp_name] = reloc
                        self.relocs.append(reloc)


    def _handle_exports(self):
        if hasattr(self._pe, 'DIRECTORY_ENTRY_EXPORT'):
            symbols = self._pe.DIRECTORY_ENTRY_EXPORT.symbols
            for exp in symbols:
                symb = WinSymbol(self, exp.name, exp.address, False, True, exp.ordinal, exp.forwarder)
                self._exports[exp.name] = symb
                self._ordinal_exports[exp.ordinal] = symb


                if exp.forwarder is not None:
                    forwardlib = exp.forwarder.split('.', 1)[0].lower() + '.dll'
                    if forwardlib not in self.deps:
                        self.deps.append(forwardlib)


    def __register_relocs(self):
        if not hasattr(self._pe, 'DIRECTORY_ENTRY_BASERELOC'):
            if self.pic:
                l.info('No relocations found in PIC binary')
            return

        for base_reloc in self._pe.DIRECTORY_ENTRY_BASERELOC:
            entry_idx = 0
            while entry_idx < len(base_reloc.entries):
                reloc_data = base_reloc.entries[entry_idx]
                if reloc_data.type == pefile.RELOCATION_TYPE['IMAGE_REL_BASED_HIGHADJ']: # special case, occupies 2 entries
                    if entry_idx == len(base_reloc.entries):
                        l.warning('PE contains corrupt base relocation table')
                        break

                    next_entry = base_reloc.entries[entry_idx]
                    entry_idx += 1
                    reloc = self._make_reloc(addr=reloc_data.rva, reloc_type=reloc_data.type, next_rva=next_entry.rva)
                else:
                    reloc = self._make_reloc(addr=reloc_data.rva, reloc_type=reloc_data.type)

                if reloc is not None:
                    self.pic = True # I've seen binaries with the DYNAMIC_BASE DllCharacteristic unset but have tons of fixup relocations
                    self.relocs.append(reloc)

                entry_idx += 1

        return self.relocs

    def _make_reloc(self, addr, reloc_type, symbol=None, next_rva=None, resolvewith=None):

        # Handle special cases first
{'fraction_non_alphanumeric': 0.09481349383266459, 'fraction_numerical': 0.005349977708426215, 'mean_word_length': 2.1534676663542642, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3474549', 'n_tokens_mistral': 4240, 'n_tokens_neox': 4026, 'n_words': 878}