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

Unpacks a UPX-packed PE file using speakeasy emulator, by breaking & dumping on section hop. Not really tested.

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 MalcatUpxUnpacker(speakeasy.Speakeasy):
    '''
    Generic UPX unpacker class
    '''

    def __init__(self, debug=False):
        super(MalcatUpxUnpacker, self).__init__(debug=debug)
        self.base_addr = 0

    def set_dump_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):
        if self.end_addr >= addr >= self.start_addr:
            print('[*] Section hop signature hit, stopping')
            self.stop()
        return True


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

unpacker = MalcatUpxUnpacker()

# Load the module
module = unpacker.load_module(data=analysis.file[:])
base = module.get_base()

print('[*] Unpacking module with section hop')
# Get the section info for "UPX0" to detect the section hop
upx0 = module.get_section_by_name('UPX0')
if not upx0:
    raise ValueError("Could not find a section named UPX0")

start = base + upx0.VirtualAddress
end = start + upx0.Misc_VirtualSize

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

# Emulate the module
unpacker.run_module(module)
print("[*] Unpacking done")
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()), "UPX.unpacked")
