import os
import subprocess

from malwareconfig.yarascanner import YaraScanner
from malwareconfig.common import PreProcessor
from tempfile import NamedTemporaryFile


class UPX(PreProcessor):
    preprocessor_name = "UPX"
    preprocessor__version = 1
    preprocessor_author = "@kevthehermit"
    preprocessor_description = "Unpack UPX Samples"

    def __init__(self):
        self.config = {}

    def pre_process(self):
        '''
        This is the main entry
        :return:
        '''

        file_data = self.file_info.file_data

        f = NamedTemporaryFile(delete=False)
        f.write(file_data)
        f.close()

        try:
            # Run UPX to unpack on the command line
            with open(os.devnull, 'w') as devnull:
                # this will also redirect stderr to /dev/null as well
                subprocess.run(['upx', '-d', f.name], stdout=devnull, stderr=devnull)

            # Read the new file
            new_data = open(f.name, 'rb').read()

            # Remove the UPX temp file
            os.unlink(f.name)

            # Set the file data to the new data
            self.file_info.file_data = new_data

            # Scan with yara
            print("    [-] Rescanning File")
            scanner = YaraScanner()
            scanner.yara_scan(self.file_info.file_data)
            self.file_info.malware_name = scanner.rule_list[0]
            print("    [-] Detected: {0}".format(self.file_info.malware_name))

        except Exception as e:
            print("Failed to Unpack sample: {0}".format(e))