generate_mesos.py

import sys, subprocess, os, zipfile, threading, time, struct

# Maximum number of processing threads to use
# idat64.exe fails silently on OOMs and stuff so we just keep this reasonable
MAX_JOBS = 4

# Name of the input zip file created by `offline_meso.ps1`
PREP_ZIP = "meso_deps.zip"

# Name of IDA executable
IDA_NAME = "idat64.exe"

# Name of the folder to write the mesos to
CACHE_FOLDER_NAME = "cache"

def usage():
    print("Usage:")
    print("    generate_mesos.py process_ida")
    print("        Processes all files in the meso_deps.zip file\n")

    print("    generate_mesos.py process_ida_whitelist <str 1> <str 2> <str ...>")
    print("        Processes files only containing one of the strings provided\n")

    print("    generate_mesos.py process_ida_blacklist <str 1> <str 2> <str ...>")
    print("        Processes files all files except for those containing one of the provided strings\n")

    print("Examples:\n")
    print("    python generate_mesos.py process_ida_whitelist system32")
    print("        Only processes files in `system32`\n")
    print("    python generate_mesos.py process_ida_blacklist ntdll.dll")
    print("        Process all files except for `ntdll.dll`\n")

    print("Path requirements for process_ida_*: must have `idat64.exe` in your PATH")
    quit()

# Make sure ida is installed and working
def ida_probe():
    PROBENAME = os.path.join("mesogen_scripts", ".idaprobe")

    def ida_error():
        # Validate it worked!
        assert os.path.exists(PROBENAME), \
            "idat64.exe is not in your PATH or it's not functioning. \
             Perhaps you need to accept a license agreement"

    # Remove probe file
    try:
        os.unlink(PROBENAME)
    except FileNotFoundError:
        pass

    # Invoke IDA to generate the idaprobe file
    try:
        subprocess.check_call([IDA_NAME, "-t", "-A",
            "-Smesogen_scripts/ida_detect.py"], shell=True)
    except:
        pass

    # Display error if file did not get created
    ida_error()

    # Remove file
    os.unlink(PROBENAME)

def process_ida(orig_name, cache_fn, cache_fn_bin, contents):
    if not os.path.exists(cache_fn):
        # Make the hirearchy for the cache file
        try:
            os.makedirs(os.path.dirname(cache_fn))
        except FileExistsError:
            pass

        # Save file to disk
        with open(cache_fn_bin, "wb") as fd:
            fd.write(contents)

        # Invoke IDA to generate the meso file
        subprocess.check_call([
            IDA_NAME, "-o%s.idb" % cache_fn, "-A",
            "-Smesogen_scripts/ida.py cmdline \"%s\" \"%s\"" % \
                    (cache_fn, os.path.basename(orig_name)),
            "-c", cache_fn_bin], shell=True)

def process(whitelist, blacklist):
    # Open the zip file generated by an offline meso script
    tf = zipfile.ZipFile(PREP_ZIP, "r")

    for member in tf.infolist():
        # If there's no file size (it's a directory) skip it
        if member.file_size == 0:
            continue

        # Check if the blacklist excludes this file
        if blacklist is not None:
            in_blacklist = filter(lambda x: x in member.filename, blacklist)
            if len(list(in_blacklist)) > 0:
                continue

        # Check if the whitelist includes a file
        if whitelist is not None:
            in_whitelist = filter(lambda x: x in member.filename, whitelist)
            if len(list(in_whitelist)) == 0:
                continue
        
        print("Processing %s" % member.filename)

        # Read the file from the archive
        contents = None
        with tf.open(member, "r") as fd:
            contents = fd.read()

        # Parse out the TimeDateStamp and SizeOfImage from the PE header
        assert contents[:2] == b"MZ"
        pe_ptr = struct.unpack("<I", contents[0x3c:0x40])[0]
        assert contents[pe_ptr:pe_ptr+4] == b"PE\0\0"
        time_date_stamp = \
                struct.unpack("<I", contents[pe_ptr+8:pe_ptr+0xc])[0]
        size_of_image = \
                struct.unpack("<I", contents[pe_ptr+0x50:pe_ptr+0x54])[0]

        image_name = os.path.split(member.filename)[1]

        # Compute meso and binary names for the cache folder
        cache_fn = os.path.join(CACHE_FOLDER_NAME, "%s_%x_%x.meso" % \
                (image_name, time_date_stamp, size_of_image))
        cache_fn_bin = os.path.join(CACHE_FOLDER_NAME, "%s_%x_%x" % \
                (image_name, time_date_stamp, size_of_image))

        # Limit number of active threads
        while threading.active_count() > MAX_JOBS:
            time.sleep(0.1)

        # Create thread
        threading.Timer(0.0, process_ida, \
                args=[image_name, cache_fn, cache_fn_bin, contents]) \
                .start()

    # Wait for all jobs to finish
    while threading.active_count() > 1:
        time.sleep(0.1)

# Check that IDA is installed and working
ida_probe()

args = sys.argv[1:]

# Super secret options
while len(args) >= 2:
    if args[0] == "--zipfile":
        # Use a custom zipfile as input
        PREP_ZIP = args[1]
        args = args[2:]
        continue
    elif args[0] == "--threads":
        # Change the number of worker jobs for meso generation
        MAX_JOBS = int(args[1])
        args = args[2:]
        continue
    break

if len(args) == 1 and args[0] == "process_ida":
    process(None, None)
elif len(args) >= 2 and args[0] == "process_ida_whitelist":
    process(args[1:], None)
elif len(args) >= 2 and args[0] == "process_ida_blacklist":
    process(None, args[1:])
else:
    usage()