"""
name: save as new ELF
category: elf
author: malcat
icon: wxART_FILETYPE_ARCHIVE

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

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

If at least one function is present in the current selection/file, the entry point of this new ELF 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

SEGMENT_ALIGN = 0x1000
IMAGEBASE = 0x8000000

ELF_HEADER_32 ="7F454C460101010000565656560000000200{machine}01000000{entrypoint}340000000000000000000000340020000100000000000000"
ELF_HEADER_64 ="7F454C460201010000565656560000000200{machine}01000000{entrypoint}000000004000000000000000000000000000000000000000400038000100000000000000"

SEGMENT_32 = "0100000000000000{section_va}{section_va}{section_rawsize}{section_vsize}07000000{align}"
SEGMENT_64 = "01000000070000000000000000000000{section_va}{section_va}{section_rawsize}{section_vsize}{align}"


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 = [ELF_HEADER_32, SEGMENT_32]
    machine = 0x3
    va_fmt = "<I"
    size_of_headers = 0x54
elif analysis.architecture == malcat.Architecture.X64:
    patterns = [ELF_HEADER_64, SEGMENT_64]
    machine = 0x3e
    va_fmt = "<Q"
    size_of_headers = 0x78
else:
    raise ValueError(f"Unsupported architecture {analysis.architecture}")


disk_size = len(data) + size_of_headers
virtual_size = align(disk_size, SEGMENT_ALIGN)

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

ctx = {
    "entrypoint": struct.pack("<I", entry_point).hex().upper(),
    "machine": struct.pack("<H", machine).hex().upper(),
    "section_va": struct.pack(va_fmt, IMAGEBASE).hex().upper(),
    "section_rawsize": struct.pack(va_fmt, disk_size).hex().upper(),
    "section_vsize": struct.pack(va_fmt, virtual_size).hex().upper(),
    "align": struct.pack(va_fmt, SEGMENT_ALIGN).hex().upper(),
}
header = "".join(patterns).format(**ctx)
elffile = bytes.fromhex(header)
elffile += data
gui.open_after(elffile, "fake_elf.bin")
