import pefile

class Binary:

    def __init__(self, file):
        self.file = file
        self.pe   = pefile.PE(file)
    @staticmethod
    def __help__():
        print("""
            Binary is a wrapper on top of pefile; static analysis utility
            Exposed methods:
                dataDir -> returns the data directories names, used by the OS for example the import table
                and export table to define the imports and exports of a portable executable
                exportTab and importTab -> returns the exported functions from a DLL and the imported DLLs
                of an executable.
                importTabFuncs -> returns a dictionary of the imported DLLs and the included exported functions
                sections -> returns a dictionary of the available sections in the executable, with additional data
        """)

    def dataDir(self):
        data_dirs = []
        for data_directory in self.pe.OPTIONAL_HEADER.DATA_DIRECTORY:
            data_dirs.append(data_directory.name)
        return data_dirs
    def exportTab(self):
        funcs = []
        for exp in pe.DIRECTORY_ENTRY_EXPORT.symbols:
            funcs.append([hex(pe.OPTIONAL_HEADER.ImageBase + exp.address),exp.name.decode('utf-8')])
        return funcs

    def importTab(self):
        imported = []
        for entry in self.pe.DIRECTORY_ENTRY_IMPORT:
            imported.append(entry.dll.decode('utf-8'))
        return imported
    def importTabFuncs(self):
        """
        listing symobls ; import table
        hash table:
            dll => [function names]
            useful to determine the capabilities of a malware sample
        """
        funcs = {}
        for entry in self.pe.DIRECTORY_ENTRY_IMPORT:
            dll_name = entry.dll.decode('utf-8')
            if dll_name not in funcs.keys(): funcs[dll_name] = []
            funcs[dll_name] = []
            for func in entry.imports:
                tmp = []
                tmp.append(func.name.decode('utf-8'))
                tmp.append(func.address)
                funcs[dll_name].append(tmp)
        return funcs
    def sections(self):
        sections = {}
        for section in self.pe.sections:
            sections[section.Name.decode('utf-8').strip("\x00")] = [hex(section.VirtualAddress), hex(section.Misc_VirtualSize), hex(section.SizeOfRawData)]
        return sections