import source files, forgetting about 938 prior commits

scripts/testbed/freeze_features.py

@@ -0,0 +1,102 @@
+'''
+Freeze capa features.
+
+Example usage:
+  freeze_features.py <test files dir>
+  freeze_features.py samples\benign
+'''
+
+import os
+import sys
+import time
+import logging
+
+import argparse
+
+from scripts.testbed import FREEZE_EXTENSION
+from capa.features.freeze import main as freeze_features
+
+
+# only process files with these extensions
+TARGET_EXTENSIONS = [
+    '.mal_',
+    '.exe_',
+    '.dll_',
+    '.sys_'
+]
+
+
+logger = logging.getLogger('check_rule')
+
+
+def freeze(input_path, reprocess):
+    if not os.path.exists(input_path):
+        raise IOError('%s does not exist or cannot be accessed' % input_path)
+
+    if os.path.isfile(input_path):
+        outfile = '%s%s' % (input_path, FREEZE_EXTENSION)
+        freeze_file(input_path, outfile, reprocess)
+
+    elif os.path.isdir(input_path):
+        logger.info('freezing features of %s files in %s', '|'.join(TARGET_EXTENSIONS), input_path)
+        for root, dirs, files in os.walk(input_path):
+            for file in files:
+                if not os.path.splitext(file)[1] in TARGET_EXTENSIONS:
+                    logger.debug('skipping non-target file: %s', file)
+                    continue
+                path = os.path.join(root, file)
+                outfile = '%s%s' % (path, FREEZE_EXTENSION)
+                freeze_file(path, outfile, reprocess)
+
+
+def freeze_file(path, output, reprocess=False):
+    logger.info('freezing features of %s', path)
+
+    if os.path.exists(output) and not reprocess:
+        logger.info('%s already exists, provide -r argument to reprocess', output)
+        return
+
+    try:
+        freeze_features([path, output])  # args: sample, output
+    except Exception as e:
+        logger.error('could not freeze features for %s: %s', path, str(e))
+
+
+def main(argv=None):
+    if argv is None:
+        argv = sys.argv[1:]
+
+    parser = argparse.ArgumentParser(description="Freeze capa features of a file or of files in a directory")
+    parser.add_argument("file_path", type=str,
+                        help="Path to file or directory 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")
+    parser.add_argument("-q", "--quiet", action="store_true",
+                        help="Disable all output but errors")
+    args = parser.parse_args(args=argv)
+
+    if args.quiet:
+        logging.basicConfig(level=logging.ERROR)
+        logging.getLogger().setLevel(logging.ERROR)
+    elif args.verbose:
+        logging.basicConfig(level=logging.DEBUG)
+        logging.getLogger().setLevel(logging.DEBUG)
+    else:
+        logging.basicConfig(level=logging.INFO)
+        logging.getLogger().setLevel(logging.INFO)
+
+    time0 = time.time()
+    try:
+        freeze(args.file_path, args.reprocess)
+    except IOError as e:
+        logger.error('%s', str(e))
+        return -1
+
+    logger.info('freezing features took %d seconds', time.time() - time0)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())