mal_unpack_runner.py

#!/usr/bin/env python3

import sys, os, subprocess
import pefile
from pathlib import Path

def mal_unp_res_to_str(returncode):
    if returncode == (-1):
        return "ERROR"
    if returncode == 0:
        return "INFO"
    if returncode == 1:
        return "NOT_DETECTED"     
    if returncode == 2:
        return "DETECTED"       
    return hex(returncode)

def get_config(sample):
    null_config = (None, None)
    try:
        pe = pefile.PE(sample, fast_load=True)
        if pe is None:
            return null_config
        is_dll = False
        if (pe.is_exe() == False):
            is_dll = True
        is_64b = None
        if (pe.OPTIONAL_HEADER.Magic & 0x10b):
            is_64b = False
        elif (pe.OPTIONAL_HEADER.Magic & 0x20b):
            is_64b = True     
        pe.close()
        if (is_64b == None):
            return null_config   
        return (is_64b, is_dll)
    except:
        return null_config   
        
def rename_sample(sample, is_dll):
    p = Path(sample)
    ext = p.suffix
    print ("Name: " +  p.stem)
    print ("Ext: " +  ext)
    
    print("Is DLL: " + str(is_dll))
    new_ext = ".exe"
    if (is_dll):
        new_ext = ".dll"
    new_name = p.stem + new_ext
    directory = str(p.parent)
    print ("Dir: " +  directory)
    p.rename(Path(p.parent, new_name))
    abs_name = directory + os.path.sep + new_name 
    print("New name: " + abs_name)
    return abs_name
    
def run_and_dump(sample, dump_dir, timeout):
    is_64b, is_dll = get_config(sample)
    if is_64b == None:
        print("[-] Not a valid PE")
        return

    print("Is 64b: " + str(is_64b))
    print("Is DLL: " + str(is_dll))
    
    orig_name = sample
    sample = rename_sample(sample, is_dll)
        
    cmd = ['mal_unpack.exe',
    '/timeout' , str(timeout),
    '/dir', dump_dir,
    '/img', sample]
    
    exe_name = sample
    if is_dll:
        cmd.append('/cmd')
        cmd.append(sample)
        if is_64b:
            exe_name = "dll_load64.exe"
        else:
            exe_name = "dll_load32.exe"

    cmd.append('/exe')
    cmd.append(exe_name)
    result = subprocess.run(cmd, check=False, capture_output=True)
    if (result.returncode is None):
        print("mal_unpack failed to run")
    else:
        print("mal_unpack result: " + mal_unp_res_to_str(result.returncode))

def main():
    if (len(sys.argv) < 2):
        print("Arg: filename")
        return
    print("starting...")
    timeout = 1000
    run_and_dump(sys.argv[1], "dumps1234", timeout)
    
if __name__ == "__main__":
    main()