olevba

olevba.py

import os
import time
import argparse
from pprint import pprint as pp
from pprint import pformat as pf
from contextlib import redirect_stdout

import oletools.oleid
from oletools.olevba import VBA_Parser, VBA_Parser_CLI, TYPE_OLE, TYPE_OpenXML, TYPE_Word2003_XML, TYPE_MHTML

def NanoToMills(nanoseconds):
    return round(nanoseconds/1000000)

def NanoToSec(nanoseconds):
    return MillsToSec(NanoToMills(nanoseconds))

def MillsToSec(milliseconds):
    return ((milliseconds/1000) % 60)

def print_perf(proc_time, perf_time):
    print("[*]: {} counter: {:>15.2f} {:3s}".format("Proc/CPU", proc_time, 'ns'))
    print("[*]: {} counter: {:>15.2f} {:3s}".format("Proc/CPU", NanoToMills(proc_time), 'ms'))
    print("[*]: {} counter: {:>15.2f} {:3s}".format("Proc/CPU", NanoToSec(proc_time), 's'))
    print()
    print("[*]: {} counter: {:>15.2f} {:3s}".format("Perf/Wall", perf_time, 'ns'))
    print("[*]: {} counter: {:>15.2f} {:3s}".format("Perf/Wall", NanoToMills(perf_time), 'ms'))
    print("[*]: {} counter: {:>15.2f} {:3s}".format("Perf/Wall", NanoToSec(perf_time), 's'))
    print()
    print()

def doprocess(filename, process_opts, analysis_opts, outfile):
    print('[!]: %s(%s, %s, %s, %s)' % ("doprocess", filename, process_opts, analysis_opts, outfile))
    print('[*]: Capturing perf counters')
    sProcT = time.process_time_ns()
    sPerfT = time.perf_counter_ns()

    print('[!]: Taking over stdout:')
    with open(outfile, 'w') as f:
        with redirect_stdout(f):
            localparser = VBA_Parser_CLI(filename)
            localparser.process_file(**process_opts)
            localparser.run_analysis(**analysis_opts)
    print('[!]: Restored stdout:')
    elapsed_proc_time_ns = time.process_time_ns() - sProcT
    elapsed_perf_time_ns = time.perf_counter_ns() - sPerfT
    print_perf(elapsed_proc_time_ns, elapsed_perf_time_ns)

def scandir(directory=None):
    """ Return files in a directory """
    ret = []
    if directory is None:   directory = os.getcwd()
    ## Loop over files
    for filename in os.scandir(directory):
        if os.path.isfile(filename):
            oid = oletools.oleid.OleID(filename)
            indicators = oid.check()
            ## Loop over indicators
            for i in indicators:
                ## Continue for each indicator until container, then break to next file
                if i.id == "container":
                    if i.value == "OLE":
                        ret.append(filename)
                        print("[*] Found an OLE container: %s" % (filename))
                        break
                    else:
                        print("[*] File %s does not appear to be an ole file. container=(%s)" % (filename, i.value))
                        break
                else:
                    continue
    return ret

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--directory', help='Select a directory to scan')
    parser.add_argument('--speed', choices=['fast', 'medium', 'slow'], default="medium", help="""
        fast: extract vba; medium: extract vba, hex-decode strings; slow: extract vba, hex-decode strings, deobfuscate; (default: %(default)s)
    """)
    args = parser.parse_args()
    return args

def getSpeedOpts(speed):
    """
        returns a tuple for the following:

        analysis_opts = {
            "show_decoded_strings":0
            "deobfuscate":1
        }
    """
    if speed == "fast":
        return (False, False)
    elif speed == "medium":
        return (True, False)
    elif speed == "slow":
        return (True, True)
    else:
        return (None, None)

def main():
    args = parse_args()
    if args.directory is None:
        args.directory = os.getcwd()

    speed_opts = getSpeedOpts(args.speed)
    analysis_opts = {
        "show_decoded_strings":speed_opts[0],
        "deobfuscate":speed_opts[1],
    }
    process_opts = {
        "show_decoded_strings": analysis_opts["show_decoded_strings"],
        "deobfuscate": analysis_opts["deobfuscate"],
        "display_code": True,
        "hide_attributes": False,
        "show_deobfuscated_code": True,
    }
    files = scandir(args.directory)

    for f in files:
        vbaparser = VBA_Parser(f)
        if not vbaparser.detect_vba_macros():
            print('[*] No VBA Macros found')
        else:
            print('[*] VBA Macros found')
            basedir = os.path.dirname(f.path)
            newfile = "%s.%s" % (f.name, "vbaparser")
            target = os.path.join(basedir, newfile)
            try:
                doprocess(f, process_opts, analysis_opts, target)
            except Exception as e:
                print("[!] Caught exception %s" % str(e))

if __name__=="__main__":
    main()