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

Unpacks a PE file using Unipacker emulator, starting at entry point. Not really tested.

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

https://github.com/unipacker/unipacker
"""
import threading
import io
import sys
import tempfile
import shutil
import os
import time

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

[link]https://github.com/unipacker/unipacker[/link]
""", format=True)
    raise


def get_unpacker(sample, auto_default_unpacker=True):
    if "UPX" in analysis.sigs:
        return unipacker.unpackers.UPXUnpacker(sample)
    else:
        return unipacker.unpackers.AutomaticDefaultUnpacker(sample)


class MalcatUnpacker(UnpackerClient):

    def __init__(self):
        UnpackerClient.__init__(self)
        self.sample = Sample(analysis.file.path)
        #self.sample.unpacker = get_unpacker(self.sample)
        self.shell_event = threading.Event()
        self.engine = UnpackerEngine(self.sample)
        self.engine.register_client(self)
        self.engine.log_mem_read = False
        self.engine.log_instr = False
        self.engine.log_mem_write = True
        self.engine.log_apicalls = True
        self.sample.unpacker.section_hopping_control = False
        self.sample.unpacker.write_execute_control = True
        self.address_updated(self.sample.unpacker.startaddr)

    def emu_started(self):
        self.started = True

    def emu_paused(self):
        self.stop()

    def emu_resumed(self):
        pass

    def emu_done(self):
        self.started = False
        self.stop()

    def address_updated(self, address):
        self.address = address

    def run(self):
        self.shell_event.clear()
        threading.Thread(target=self.engine.emu).start()
        self.shell_event.wait()

    def stop(self):
        if self.started:
            self.shell_event.clear()
            self.engine.stop()
        else:
            self.shell_event.set()
        if self.started:
            self.engine.stop()
        self.shell_event.set()



##########################################

olddir = os.getcwd()
tmpdir = tempfile.mkdtemp(prefix=analysis.file.path)
try:
    os.chdir(tmpdir)
    unpacker = MalcatUnpacker()        
    unpacker.run()

    if not os.path.exists("unpacked.exe"):
        print("No unpacked file could be found")
    else:
        with open("unpacked.exe", "rb") as f:
            gui.open_after(f.read(), "unpacked.exe")

finally:
    os.chdir(olddir)
    shutil.rmtree(tmpdir)

