import pefile
import os
import json

def GetExportInformation(FileName: str):

    try:
        directory = [pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_EXPORT"]]
        file_content = pefile.PE(FileName, fast_load=True)
        file_content.parse_data_directories(directories=directory)
        exports = [e.name.decode() for e in file_content.DIRECTORY_ENTRY_EXPORT.symbols if e.name]
        return sorted(exports)
    except Exception as e:
        return []

def GetImportInformation(FileName: str):
    try:
        file_content = pefile.PE(FileName, fast_load=True)
        file_content.parse_data_directories()
        imports = []
        for entry in file_content.DIRECTORY_ENTRY_IMPORT:
            dll_name = entry.dll.decode() if entry.dll else "Unknown"
            functions = [imp.name.decode() for imp in entry.imports if imp.name]
            imports.append({"dll": dll_name, "functions": functions})
        return imports
    except Exception as e:
        return []

if __name__ == "__main__":
    results = {}
    file_names = os.listdir()

    for file in file_names:
        try:
            print(f"Processing file: {file}")
            imports = GetImportInformation(file)
            exports = GetExportInformation(file)
            results[file] = {
                "imports": imports,
                "exports": exports
            }
        except Exception as e:
            print(f"Failed to process {file}: {str(e)}")
            results[file] = {
                "error": "Failed to process this file."
            }

    # Save the results to a JSON file
    with open("pe_info.json", "w") as json_file:
        json.dump(results, json_file, indent=4)

    print("Output written to pe_info.json")