Initial commit

Author: gamozolabs

.gitignore

@@ -0,0 +1,9 @@
+/target
+**/*.rs.bk
+*.i64
+*.idb
+*.meso
+cache
+coverage.txt
+*.zip
+

Cargo.lock

@@ -0,0 +1,11 @@
+[package]
+name = "mesos"
+version = "0.1.0"
+authors = ["bfalk"]
+
+[dependencies]
+winapi = { version = "0.3.5", features = ["debugapi", "winbase", "memoryapi", "processthreadsapi", "errhandlingapi", "handleapi", "securitybaseapi", "consoleapi", "winerror", "wow64apiset"] }
+
+[profile.release]
+debug = true
+

Cargo.toml

@@ -0,0 +1,54 @@
+import collections, re, math
+
+fft = re.compile("[0-9a-f]{16} \| Freq: +([0-9]+) \| +(.*?)\+0x([0-9a-f]+) \| (.*?)\n")
+
+image_base  = idaapi.get_imagebase()
+ida_modname = GetInputFile().lower()
+
+SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+
+inp = open(os.path.join(SCRIPT_DIR, "..", "coverage.txt"), "r").read()
+
+# Reset all coloring in all functions
+for funcea in Functions():
+    f  = idaapi.get_func(funcea)
+    fc = idaapi.FlowChart(f)
+
+    for block in fc:
+        ea = block.startEA
+        while ea <= block.endEA:
+            set_color(ea, CIC_ITEM, DEFCOLOR)
+            ea = idc.NextHead(ea)
+
+freqs = collections.Counter()
+
+# Parse input coverage file
+for thing in fft.finditer(inp):
+    freq   = int(thing.group(1), 10)
+    module = thing.group(2)
+    offset = int(thing.group(3), 16)
+
+    # Skip non-matching modules
+    if module.lower() != ida_modname:
+        continue
+
+    freqs[image_base + offset] = freq
+
+# Apply coloring
+for addr, freq in freqs.most_common()[::-1]:
+    function_addr = get_func_attr(addr, FUNCATTR_START)
+    func_entry_freq = freqs[function_addr]
+
+    if func_entry_freq == 0:
+        func_entry_freq = 1
+
+    # Log value between [0.0, 1.0)
+    dist = math.log(float(freq) / float(func_entry_freq) + 1.0)
+    dist = min(dist, float(1.0))
+
+    color = 0x808080 + (int((1 - dist) * 100.0) << 8) + (int(dist * 100.0) << 0)
+    print("%10d | 0x%.16x | %s" % (freq, addr, get_func_off_str(addr)))
+
+    set_color(addr, CIC_ITEM, color)
+    set_cmt(addr, "Freq: %d | Func entry: %.2f" % (freq, float(freq) / float(func_entry_freq)), False)
+

coverage_scripts/ida.py

@@ -0,0 +1,120 @@
+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"
+
+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()
+
+def process_job(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)
+
+        try:
+            # 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)
+        except FileNotFoundError:
+            print("Could not find idat64.exe, is it not in your path?")
+            quit(-1)
+
+def process_ida(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]
+
+        # Compute meso and binary names for the cache folder
+        cache_fn = "%s_%x_%x.meso" % \
+                (member.filename, time_date_stamp, size_of_image)
+        cache_fn_bin = "%s_%x_%x" % \
+                (member.filename, 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_job, \
+                args=[member.filename, cache_fn, cache_fn_bin, contents]) \
+                .start()
+
+    # Wait for all jobs to finish
+    while threading.active_count() > 1:
+        time.sleep(0.1)
+
+if len(sys.argv) == 2 and sys.argv[1] == "process_ida":
+    process_ida(None, None)
+elif len(sys.argv) >= 2 and sys.argv[1] == "process_ida_whitelist":
+    process_ida(sys.argv[2:], None)
+elif len(sys.argv) >= 2 and sys.argv[1] == "process_ida_blacklist":
+    process_ida(None, sys.argv[2:])
+else:
+    usage()
+

generate_mesos.py

@@ -0,0 +1,44 @@
+from idautils import *
+from idaapi import *
+from idc import *
+
+# Wait for auto analysis to complete
+idaapi.auto_wait()
+
+print("Analysis done, generating meso")
+
+image_base = idaapi.get_imagebase()
+
+output = bytearray()
+
+SEPERATOR = "~~ed6ed28d321bbdc8~~"
+
+input_name = GetInputFile()
+
+if len(ARGV) >= 4 and ARGV[1] == "cmdline":
+    input_name = ARGV[3]
+
+for funcea in Functions():
+    f  = idaapi.get_func(funcea)
+    fc = idaapi.FlowChart(f)
+
+    for block in fc:
+        if is_code(getFlags(block.startEA)):
+            output += "%s%s%s%s%x%s%s\n" % \
+                    ("single", SEPERATOR,
+                     input_name, SEPERATOR,
+                     block.startEA - image_base, SEPERATOR,
+                     get_func_off_str(block.startEA))
+
+filename = "%s/%s.meso" % (os.path.dirname(os.path.abspath(__file__)), input_name)
+if len(ARGV) >= 4 and ARGV[1] == "cmdline":
+    filename = ARGV[2]
+
+open(filename, "wb").write(output)
+
+print("Generated meso: %s" % filename)
+
+# Exit only if we were invoked from the command line
+if len(ARGV) >= 4 and ARGV[1] == "cmdline":
+    idc.Exit(0)
+

mesogen_scripts/ida.py

@@ -0,0 +1,85 @@
+Param (
+    [Parameter(Mandatory=$true)][string]$TargetPid,
+    [string]$OutputZip = "meso_deps.zip"
+)
+
+function is64bit($a) {
+    try {
+        Add-Type -MemberDefinition @'
+[DllImport("kernel32.dll", SetLastError = true, 
+ CallingConvention = CallingConvention.Winapi)]
+[return: MarshalAs(UnmanagedType.Bool)]
+public static extern bool IsWow64Process(
+ [In] System.IntPtr hProcess,
+ [Out, MarshalAs(UnmanagedType.Bool)] out bool wow64Process);
+'@ -Name NativeMethods -Namespace Kernel32
+    }
+    catch {}
+    $is32Bit = [int]0
+    if (!$a.Handle) {
+        echo "Unable to open handle for process: Does the proceses exist? Do you have adequate permissions?"
+        exit
+    }
+    if ([Kernel32.NativeMethods]::IsWow64Process($a.Handle, [ref]$is32Bit)) {
+        $(if ($is32Bit) {$false} else {$true})
+    } else {
+        "IsWow64Process call failed"
+        exit
+    }
+}
+
+# Create a new temp directory based on a random GUID
+function New-TemporaryDirectory {
+    $parent = [System.IO.Path]::GetTempPath()
+    [string] $name = [System.Guid]::NewGuid()
+    $name = "mesotmp_" + $name
+    New-Item -ItemType Directory -Path (Join-Path $parent $name)
+}
+
+$pshell_bitness = (is64bit(Get-Process -Id $PID))
+echo "Powershell is 64-bit: $pshell_bitness"
+
+$target_bitness = (is64bit(Get-Process -Id $TargetPid))
+echo "Target     is 64-bit: $target_bitness"
+
+# Validate bitnesses match
+if ($pshell_bitness -ne $target_bitness) {
+    echo "Your Powershell bitness does not match the target bitness"
+    echo "Use 32-bit Powershell for 32-bit processes or 64-bit Powershell for 64-bit processes"
+    echo "This is to get around pathing issues between things like C:\windows\system32 and C:\windows\syswow64"
+    exit
+}
+
+# Get all the module/executable paths from a running process
+$paths = (Get-Process -Id $TargetPid -Module -FileVersionInfo).FileName
+
+# Create a new temporary directory
+$dirname = New-TemporaryDirectory
+
+ForEach ($path in $paths) {
+    # Convert the path to not have a ":" in it by replacing it with a "_" and
+    # then convert it to a lowercase string
+    $lowerpath = ([string](Get-ChildItem -Path $path)).replace(":", "_").ToLower()
+
+    # Prepend a root-level folder "cache" to all paths that will go in the zip
+    $lowerpath = (Join-Path "cache" $lowerpath)
+
+    # Compute this path in the temp folder
+    $hirearchy = (Join-Path $dirname $lowerpath)
+
+    # Get the parent directory from this filename
+    $parent = (Split-Path $hirearchy -Parent)
+
+    # Create the directory if it doesn't exist
+    if (!(Test-Path -path $parent)) { New-Item -ItemType Directory -Path $parent | Out-Null }
+
+    # Copy the file :)
+    Copy-Item -Path $path -Destination $hirearchy
+}
+
+# Create zip from all the files in the temp folder
+Compress-Archive -Force -Path (Join-Path $dirname *) -DestinationPath $OutputZip
+
+# Remove temp directory
+Remove-Item -Recurse -Force $dirname
+