"""
name: Speakeasy (whole PE)
category: emulation
filetype: PE
author: malcat
icon: wxART_DISASM

Unpacks a PE file using Speakeasy emulator, starting at entry point.

Note: you need to install speakeasy and its dependancies for this script to work
On windows, make sure to check "Use system Python" in options dialog.

https://github.com/fireeye/speakeasy
"""
import argparse


try:
    import speakeasy
except ImportError:
    gui.print("""
You need to install speakeasy and its dependancies for this script to work
On windows, make sure to check "Use system Python" in options dialog.

[url]https://github.com/fireeye/speakeasy[/url]
""", format=True)
    raise


class MalcatPeUnpacker(speakeasy.Speakeasy):
    '''
    Generic PE unpacker class
    '''

    def __init__(self, debug=False, timeout=120):
        super(MalcatPeUnpacker, self).__init__(debug=debug)
        self.config["timeout"] = timeout
        self.base_addr = 0
        self.break_address = None

    def set_valid_range(self, base, start, end):
        # Set some variables used to identify when a dump should occur
        self.base_addr = base
        self.start_addr = start
        self.end_addr = end

    def code_hook(self, emu, addr, size, ctx):
        print(">{:x}".format(addr))

        if addr < 0x70000000  and (addr < self.start_addr or addr > self.end_addr):
            print('[*] Section hop signature hit at {:x} (out of {:x}-{:x} range), stopping'.format(addr, self.start_addr, self.end_addr))
            self.break_address = addr
            self.stop()
        return True


if analysis.type != "PE":
    raise ValueError("not a PE file")

unpacker = MalcatPeUnpacker(timeout=1200)

# Load the module
module = unpacker.load_module(path=None, data=analysis.file.read(0, analysis.file.size))
base = module.get_base()

start = analysis.map.base
end = start + len(analysis.map)

unpacker.set_valid_range(base, start, end)
# Add the callback
#unpacker.add_code_hook(unpacker.code_hook)

# Emulate the module
unpacker.run_module(module, all_entrypoints=True)
print("[*] Unpacking done")
for dump in unpacker.get_memory_dumps():
    # first try to open any unpacked MZ
    name, start, size, _, _, content = dump
    print(f"{name or '???':32.32s}  {start:09x}-{start+size:09x} ({size: 10d} bytes) -> {content[:32]}...")
    if unpacker.break_address is not None and start <= unpacker.break_address < start + size:
        gui.open_after(content, name)
        break
    elif name and not name.startswith("emu.") and content.startswith(b"MZ"):
        gui.open_after(content, name)
        break
else:
    import json
    print(json.dumps(unpacker.get_report()["entry_points"], indent=4))
    # open original exe
    mm = unpacker.get_address_map(unpacker.base_addr)
    # TODO: Fixup the import table after dumping
    gui.open_after(unpacker.mem_read(mm.get_base(), mm.get_size()), "PE.unpacked")
