feat: add JSON output

Author: H.L.J.Laloge

main.py

@@ -1,15 +1,26 @@
 #!/usr/bin/env python3
 
 import argparse
+import sys
+import json
 
-from polydet import magic, scan, rules
+from polydet import magic, scan, rules, PolyglotLevel
 
 
-def display_results(results: [(str, {})], indent=False):
-    for result in results.items():
-        if indent:
-            print('\t', end='')
-        print('%s: %s' % (result[0], result[1]))
+class PolyglotLevelEncoder(json.JSONEncoder):
+    def default(self, o):
+        if isinstance(o, PolyglotLevel):
+            return {
+                'isValid': o.is_valid,
+                'suspiciousChunks': o.suspicious_chunks,
+                'embedded': sorted(o.embedded)
+            }
+
+        return super().default(o)
+
+
+def display_results(results: {str: {str: PolyglotLevel}}, fp):
+    json.dump(results, fp, cls=PolyglotLevelEncoder)
 
 
 def create_arg_parser():
@@ -22,6 +33,7 @@ def create_arg_parser():
     arg_parser.add_argument('-r', '--rules', dest='rules', type=str, help='File to load and store rules to speed up the process.')
     arg_parser.add_argument('-c', '--recompile', dest='recompile', action='store_true', help='Re-compile rules. '
                                                                                              'Require --rules')
+    arg_parser.add_argument('-o', '--output', default='-', help='Output file')
     return arg_parser
 
 
@@ -39,12 +51,15 @@ def main():
         else:
             rules.load_or_compile(args.rules)
 
-    if len(args.files) == 1:
-        display_results(scan(args.files[0], use_magic=args.magic))
+    results = dict()
+    for filename in args.files:
+        results[filename] = scan(filename, use_magic=args.magic)
+
+    if args.output != '-':
+        with open(args.output) as output:
+            display_results(results, output)
     else:
-        for filename in args.files:
-            print('%s:' % filename)
-            display_results(scan(filename, use_magic=args.magic), indent=True)
+        display_results(results, sys.stdout)
 
 
 if __name__ == '__main__':

polydet/polyglot_level.py

@@ -1,6 +1,8 @@
 class PolyglotLevel:
 
-    def __init__(self, is_valid = True, suspicious_chunks: [(int, int)] = None, embedded: set = None):
+    __slots__ = ['is_valid', 'suspicious_chunks', 'embedded']
+
+    def __init__(self, is_valid=True, suspicious_chunks: [(int, int)] = None, embedded: set = None):
         self.is_valid = is_valid
         self.suspicious_chunks = suspicious_chunks if suspicious_chunks is not None else []
         self.embedded = embedded if embedded is not None else set()