mirror of
https://github.com/mandiant/capa.git
synced 2025-12-22 07:10:29 -08:00
136 lines
4.0 KiB
Python
136 lines
4.0 KiB
Python
'''
|
|
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())
|