"""
name: save as new PE
category: pe
author: malcat
icon: wxART_FILETYPE_ARCHIVE

Save the current selection (or the whole file if no selection) into a new PE file. The constructed PE file will have one RWX section that encloses all the code + data and no import.

Malcat will use the current CPU architecture to infer the PE's machine target.

If at least one function is present in the current selection/file, the entry point of this new PE file will be set to the first of these functions, otherwise it is left blank.

This script is useful to submit shellcode to sandboxes for instance. Currently only works for x86 and x64.
"""

import struct
import malcat

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

SECTION_ALIGN = 0x1000
FILE_ALIGN = 0x200
SECTION_RVA = SECTION_ALIGN
SECTION_OFFSET = FILE_ALIGN

MZ_HEADER ="4D5A90000300000004000000FFFF0000B800000000000000400000000000000000000000000000000000000000000000000000000000000000000000800000000E1FBA0E00B409CD21B8014CCD21546869732070726F6772616D2063616E6E6F742062652072756E20696E20444F53206D6F64652E0D0D0A2400000000000000"

PE_HEADER_32 = "504500004C01010009F9A4670000000000000000E0000F03"
PE_HEADER_64 = "5045000064860100F0C679680000000000000000F0002F02"

OPTHEADER_32 = "0B010E1D000000000000000000000000{entrypoint}{section_rva}0000000000004000001000000002000005000100000000000500010000000000{sizeofimage}0004000000000000020000800000100000100000000010000010000000000000100000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000074040000000000000000000000000000000000000000000000000000"
OPTHEADER_64 = "0B020E2C000000000000000000000000{entrypoint}{section_rva}0000008001000000001000000002000005000000000000000600000000000000{sizeofimage}000400000000000002000000000010000000000000100000000000000000100000000000001000000000000000000000100000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000"

SECTION = "2E726177636F6465{virtualsize}{section_rva}{sizeofrawdata}{section_offset}000000000000000000000000600000E0"


if len(analysis.sel):
    start_off = analysis.a2p(analysis.sel.start)
    delta = start_off
    end_off = start_off + len(analysis.sel)
    data = analysis.file[start_off:end_off]
else:
    delta = 0
    data = analysis.file[:]

if analysis.architecture == malcat.Architecture.X86:
    patterns = [MZ_HEADER, PE_HEADER_32, OPTHEADER_32, SECTION]
elif analysis.architecture == malcat.Architecture.X64:
    patterns = [MZ_HEADER, PE_HEADER_64, OPTHEADER_64, SECTION]
else:
    raise ValueError(f"Unsupported architecture {analysis.architecture}")


virtual_size = align(len(data), SECTION_ALIGN)
disk_size = align(len(data), FILE_ALIGN)

for fn in analysis.fns:
    if analysis.a2p(fn.start) >= delta and analysis.a2p(fn.end) <= delta + len(data):
        entry_point = SECTION_RVA + analysis.a2p(fn.start) - delta
        break
else:
    entry_point = SECTION_RVA

ctx = {
    "entrypoint": struct.pack("<I", entry_point).hex().upper(),
    "section_rva": struct.pack("<I", SECTION_RVA).hex().upper(),
    "section_offset": struct.pack("<I", SECTION_OFFSET).hex().upper(),
    "virtualsize": struct.pack("<I", virtual_size).hex().upper(),
    "sizeofrawdata": struct.pack("<I", disk_size).hex().upper(),
    "sizeofimage": struct.pack("<I", SECTION_RVA + virtual_size).hex().upper(),
}
header = "".join(patterns).format(**ctx)
pefile = bytes.fromhex(header)
if len(pefile) > SECTION_OFFSET:
    raise ValueError("header too big")
elif len(pefile) < SECTION_OFFSET:
    pefile += b"\x00" * (SECTION_OFFSET - len(pefile))
pefile += data
if (len(pefile) % FILE_ALIGN) != 0:
    pefile += b"\x00" * (FILE_ALIGN - (len(pefile) % FILE_ALIGN))
gui.open_after(pefile, "fake_pe.exe")
