"""
name: Debloat last section
category: pe
filetype: PE
author: malcat
icon: wxART_FILE_COMPRESSED

Removes all the trailing 00 bytes from the last section / overlay and patch headers as needed.

Note: this is only a tiny script for one use case of file bloating. If you want a proper debloat tool handling more use cases, try https://github.com/Squiblydoo/debloat.
"""
import re


def align(val, what, down=False):
    if val % what:
        if down:
            val -= val % what
        else:
            val += what - (val % what)
    return val


def debloat_last_section(analysis):
    file_alignement = analysis.struct["OptionalHeader"]["FileAlignment"]

    #get last section
    last = None
    for s in analysis.struct["Sections"]:
        if last is None or s["PointerToRawData"] > last["PointerToRawData"]:
            last = s
    if last is None:
        raise ValueError("no physical section")

    # read last section data
    section_data = analysis.file[last["PointerToRawData"]:last["PointerToRawData"]+last["SizeOfRawData"]]

    # get last non-null byte in last section
    reverted = section_data[::-1]
    m = re.search(rb"[^\x00]", reverted)
    if not m:
        delta_last_non_zero = len(reverted)
    else:
        delta_last_non_zero = len(reverted) - m.end(0)

    to_remove = delta_last_non_zero
    if to_remove < file_alignement:
        # nothing to remove there
        print(f"Not enough bytes to remove in last section: {to_remove} < {file_alignement}")
        return None
    to_remove = align(to_remove, file_alignement, True)

    print("found last non-zero byte at offset #{:x} in {}, removing last {} bytes in section !".format(last["PointerToRawData"] + delta_last_non_zero, last["Name"].replace("\x00", ""), to_remove))

    # fix last_section 
    last["VirtualSize"] -= to_remove
    last["SizeOfRawData"] -= to_remove
    # fix OptionalHeader
    analysis.struct["OptionalHeader"]["SizeOfImage"] -= to_remove

    start = last["PointerToRawData"] + delta_last_non_zero
    return analysis.file[:start] + analysis.file[start + to_remove:]


def debloat_overlay(analysis):
    overlay = analysis.map.regions[-1]
    overlay_data = analysis.file[overlay.phys:overlay.phys + overlay.phys_size]

    reverted = overlay_data[::-1]
    m = re.search(rb"[^\x00]", reverted)
    if not m:
        delta_last_non_zero = len(reverted)
    else:
        delta_last_non_zero = len(reverted) - m.end(0)

    to_remove = delta_last_non_zero
    if to_remove < 256:
        print(f"Not enough bytes to remove in overlay: {to_remove} < 256")
        return None
    print("found last non-zero byte at offset #{:x} in overlay: removing last {} bytes in overlay !".format(overlay.phys + delta_last_non_zero, to_remove))
    return analysis.file[:-to_remove]


if analysis.type != "PE":
    raise ValueError("File is not PE")

debloated = None
if analysis.map.regions[-1].name == "overlay":
    debloated = debloat_overlay(analysis)
else:
    debloated = debloat_last_section(analysis)


if debloated is not None:
    # open patched file
    gui.open_after(debloated, "debloated")
