Clean junk PE overlay

clean_pe_overlay.py

# Clean PE files that have a lot of junk after its end to avoid AV scanners and slow down analysis tools

import pefile
import sys
import os

TRESHOLD = 100

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print(f"{sys.argv[0]} <bloated>")
        sys.exit(1)
    

    bloated_pe_path = sys.argv[1]

    if not os.path.exists(bloated_pe_path):
        print(f"Unable to open file {bloated_pe_path}")
        sys.exit(1)

    data = b''
    with open(bloated_pe_path, 'rb') as b_fd:
        data = b_fd.read()
    

    bloated_pe = pefile.PE(data=data, fast_load=True)
    overlay_offset = bloated_pe.get_overlay_data_start_offset()

    if not overlay_offset:
        print(f"No overlay detected at {bloated_pe_path}")
        bloated_pe.close()
        exit(0)
    
    bloated_pe_size = len(data)

    diff = bloated_pe_size - overlay_offset 
    if diff < TRESHOLD:
        print("Overlay is not big enough to cause troubles, bye bye")
        sys.exit(0)

    print(f"Detected a big overlay, removing...")
    
    data = data[:overlay_offset]
    new_name = f"clean_{os.path.basename(bloated_pe_path)}"

    with open(new_name, "wb") as new_fd:
        new_fd.write(data)
    
    print(f"Saved as {new_name}")