pep8: black

scripts/testbed/start_ida_dump_fnames.py

@@ -1,4 +1,4 @@
-'''
+"""
 Start IDA Pro in autonomous mode to dump JSON file of function names { fva: fname }.
 Processes a single file or a directory.
 Only runs on files with supported file extensions.
@@ -6,7 +6,7 @@ Only runs on files with supported file extensions.
 Example usage:
   start_ida_dump_fnames.py <candidate files dir>
   start_ida_dump_fnames.py samples\benign
-'''
+"""
 
 import os
 import sys
@@ -19,18 +19,18 @@ import argparse
 
 from scripts.testbed import FNAMES_EXTENSION
 
-IDA32_PATH = 'C:\\Program Files\\IDA Pro 7.3\\ida.exe'
-IDA64_PATH = 'C:\\Program Files\\IDA Pro 7.3\\ida64.exe'
+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
-DUMP_SCRIPT_PATH = os.path.abspath('_dump_fnames.py')
+DUMP_SCRIPT_PATH = os.path.abspath("_dump_fnames.py")
 
 SUPPORTED_EXTENSIONS = [
-    '.exe_',
-    '.dll_',
-    '.sys_',
-    '.idb',
-    '.i64',
+    ".exe_",
+    ".dll_",
+    ".sys_",
+    ".idb",
+    ".i64",
 ]
 
 
@@ -38,40 +38,43 @@ logger = logging.getLogger(__name__)
 
 
 def call_ida_dump_script(sample_path, reprocess):
-    ''' call IDA in autonomous mode and return True if success, False on failure '''
-    logger.info('processing %s (MD5: %s)', sample_path, get_md5_hexdigest(sample_path))
+    """ call IDA in autonomous mode and return True if success, False on failure """
+    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':
+    if os.path.splitext(sample_path)[-1] == ".i64":
         IDA_PATH = IDA64_PATH
     else:
         IDA_PATH = IDA32_PATH
 
-    if sample_path.endswith('.idb') or sample_path.endswith('.i64'):
+    if sample_path.endswith(".idb") or sample_path.endswith(".i64"):
         sample_path = sample_path[:-4]
 
-    fnames = '%s%s' % (sample_path, FNAMES_EXTENSION)
+    fnames = "%s%s" % (sample_path, FNAMES_EXTENSION)
     if os.path.exists(fnames) and not reprocess:
-        logger.info('%s already exists and contains %d function names, provide -r argument to reprocess',
-                    fnames, len(get_function_names(fnames)))
+        logger.info(
+            "%s already exists and contains %d function names, provide -r argument to reprocess",
+            fnames,
+            len(get_function_names(fnames)),
+        )
         return True
 
     out_path = os.path.split(fnames)[-1]  # relative to IDA database file
-    args = [IDA_PATH, '-A', '-S%s "%s"' % (DUMP_SCRIPT_PATH, out_path), sample_path]
-    logger.debug('calling "%s"' % ' '.join(args))
+    args = [IDA_PATH, "-A", '-S%s "%s"' % (DUMP_SCRIPT_PATH, out_path), sample_path]
+    logger.debug('calling "%s"' % " ".join(args))
     subprocess.call(args)
 
     if not os.path.exists(fnames):
-        logger.warning('%s was not created', fnames)
+        logger.warning("%s was not created", fnames)
         return False
 
-    logger.debug('extracted %d function names to %s', len(get_function_names(fnames)), fnames)
+    logger.debug("extracted %d function names to %s", len(get_function_names(fnames)), fnames)
     return True
 
 
 def get_md5_hexdigest(sample_path):
     m = hashlib.md5()
-    with open(sample_path, 'rb') as f:
+    with open(sample_path, "rb") as f:
         m.update(f.read())
     return m.hexdigest()
 
@@ -79,19 +82,17 @@ def get_md5_hexdigest(sample_path):
 def get_function_names(fnames_file):
     if not os.path.exists(fnames_file):
         return None
-    with open(fnames_file, 'r') as f:
+    with open(fnames_file, "r") as f:
         return json.load(f)
 
 
 def main():
     parser = argparse.ArgumentParser(
-        description='Launch IDA Pro in autonomous mode to dump function names of a file or of files in a directory')
-    parser.add_argument('file_path', type=str,
-                        help='File or directory path to analyze')
-    parser.add_argument('-r', '--reprocess', action='store_true', default=False,
-                        help='Overwrite existing analysis')
-    parser.add_argument('-v', '--verbose', action='store_true',
-                        help='Enable verbose output')
+        description="Launch IDA Pro in autonomous mode to dump function names of a file or of files in a directory"
+    )
+    parser.add_argument("file_path", type=str, help="File or directory path to analyze")
+    parser.add_argument("-r", "--reprocess", action="store_true", default=False, help="Overwrite existing analysis")
+    parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose output")
     args = parser.parse_args(args=sys.argv[1:])
 
     if args.verbose:
@@ -102,7 +103,7 @@ def main():
         logging.getLogger().setLevel(logging.INFO)
 
     if not os.path.exists(args.file_path):
-        logger.warning('%s does not exist', args.file_path)
+        logger.warning("%s does not exist", args.file_path)
         return -1
 
     if os.path.isfile(args.file_path):
@@ -111,21 +112,21 @@ def main():
 
     errors = 0
 
-    logger.info('processing files in %s with file extension %s', args.file_path, '|'.join(SUPPORTED_EXTENSIONS))
+    logger.info("processing files in %s with file extension %s", args.file_path, "|".join(SUPPORTED_EXTENSIONS))
     for root, dirs, files in os.walk(args.file_path):
         for file in files:
             if not os.path.splitext(file)[1] in SUPPORTED_EXTENSIONS:
-                logger.debug('%s does not have supported file extension', file)
+                logger.debug("%s does not have supported file extension", file)
                 continue
             path = os.path.join(root, file)
             if not call_ida_dump_script(path, args.reprocess):
                 errors += 1
 
     if errors:
-        logger.warning('encountered %d errors', errors)
+        logger.warning("encountered %d errors", errors)
 
     return 0
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     sys.exit(main())