import source files, forgetting about 938 prior commits

scripts/testbed/start_ida_export_fimages.py

@@ -0,0 +1,135 @@
+'''
+Start IDA Pro in autonomous mode to export images of function graphs.
+
+Example usage:
+  start_ida_export_fimages.py <target file> <output dir> -f <function list>
+  start_ida_export_fimages.py test.exe imgs -f 0x401000,0x402F90
+'''
+
+import os
+import imp
+import sys
+import hashlib
+import logging
+import subprocess
+
+import argparse
+
+try:
+    imp.find_module('graphviz')
+    from graphviz import Source
+    graphviz_found = True
+except ImportError:
+    graphviz_found = False
+
+
+IDA32_PATH = 'C:\\Program Files\\IDA Pro 7.3\\ida.exe'
+IDA64_PATH = 'C:\\Program Files\\IDA Pro 7.3\\ida64.exe'
+
+# expected in same directory as this file
+EXPORT_SCRIPT_PATH = os.path.abspath('_export_fimages.py')
+
+
+logger = logging.getLogger(__name__)
+
+
+def export_fimages(file_path, out_dir, functions, manual=False):
+    '''
+    Export images of function graphs.
+    :param file_path: file to analyze
+    :param out_dir: output directory
+    :param functions: list of strings of hex formatted fvas
+    :param manual: non-autonomous mode
+    :return: True on success, False otherwise
+    '''
+    if not graphviz_found:
+        logger.warning('please install graphviz to export images')
+        return False
+
+    if not os.path.exists(out_dir):
+        os.mkdir(out_dir)
+
+    script_args = [os.path.abspath(out_dir)] + functions
+    call_ida_script(EXPORT_SCRIPT_PATH, script_args, file_path, manual)
+
+    img_count = 0
+    for root, dirs, files in os.walk(out_dir):
+        for file in files:
+            if not file.endswith('.dot'):
+                continue
+            try:
+                s = Source.from_file(file, directory=out_dir)
+                s.render(file, directory=out_dir, format='png', cleanup=True)
+                img_count += 1
+            except BaseException:
+                logger.warning('graphviz error rendering file')
+    if img_count > 0:
+        logger.info('exported %d function graph images to "%s"', img_count, os.path.abspath(out_dir))
+        return True
+    else:
+        logger.warning('failed to export function graph images')
+        return False
+
+
+def call_ida_script(script_path, script_args, sample_path, manual):
+    logger.info('processing %s (MD5: %s)', sample_path, get_md5_hexdigest(sample_path))
+
+    # TODO detect 64-bit binaries
+    if os.path.splitext(sample_path)[-1] == '.i64':
+        IDA_PATH = IDA64_PATH
+    else:
+        IDA_PATH = IDA32_PATH
+
+    args = [IDA_PATH, '-A', '-S%s %s' % (script_path, ' '.join(script_args)), sample_path]
+
+    if manual:
+        args.remove('-A')
+
+    logger.debug('calling "%s"' % ' '.join(args))
+    if subprocess.call(args) == 0:
+        return True
+    else:
+        return False
+
+
+def get_md5_hexdigest(sample_path):
+    m = hashlib.md5()
+    with open(sample_path, 'rb') as f:
+        m.update(f.read())
+    return m.hexdigest()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Launch IDA Pro in autonomous mode to export images of function graphs")
+    parser.add_argument("file_path", type=str,
+                        help="File to export from")
+    parser.add_argument("out_dir", type=str,
+                        help="Export target directory")
+    parser.add_argument("-f", "--functions", action="store",
+                        help="Comma separated list of functions to export")
+    parser.add_argument("-m", "--manual", action="store_true",
+                        help="Manual mode: show IDA dialog boxes")
+    parser.add_argument("-v", "--verbose", action="store_true",
+                        help="Enable verbose output")
+    args = parser.parse_args(args=sys.argv[1:])
+
+    if args.verbose:
+        logging.basicConfig(level=logging.DEBUG)
+        logging.getLogger().setLevel(logging.DEBUG)
+    else:
+        logging.basicConfig(level=logging.INFO)
+        logging.getLogger().setLevel(logging.INFO)
+
+    if not os.path.isfile(args.file_path):
+        logger.warning('%s is not a file', args.file_path)
+        return -1
+
+    functions = args.functions.split(',')
+    export_fimages(args.file_path, args.out_dir, functions, args.manual)
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())