Adjusted CLI

1 files changed, 70 insertions, 6 deletions

diff --git a/yara-compiler.py b/yara-compiler.py
index 4e0c77d..8050b4b 100644
--- a/yara-compiler.py
+++ b/yara-compiler.py
@@ -3,11 +3,47 @@ import os
 import re
 from yara import *
 
+class YACFormat(object):
+    @staticmethod
+    def write(filename, data):
+        raise NotImplementedError
+
+class YACFormatBinary(YACFormat):
+    STR = "binary"
+    EXT = "yac"
+
+    @staticmethod
+    def write(filename, data):
+        with open(filename, "wb") as f:
+            f.write(data)
+
+class YACFormatCHeader(YACFormat):
+    STR = "c_header"
+    EXT = "h"
+    MODE = "w"
+
+    @staticmethod
+    def write(filename, data):
+        with open(filename, "w") as f:
+            f.write("#ifndef FILE_{}_H\n#define FILE_{}_H\n\n#include <stdint.h>\n\n".format(os.path.splitext(os.path.basename(filename))[0], os.path.splitext(os.path.basename(filename))[0]))
+            f.write("{} FILE_{}[{}] = {{\n\t".format("uint8_t", os.path.splitext(os.path.basename(filename))[0], len(data)))
+            i = 0
+            for d in data:
+                i += 1
+                f.write("{}{}".format(d, ",\n\t" if (i % 16 == 0) else ", "))
+            f.write("\n};\n\n#endif\n")
+
+DICT_STR_YACFORMAT = {
+    YACFormatBinary.STR : YACFormatBinary,
+    YACFormatCHeader.STR : YACFormatCHeader
+}
+
 def dir_path(string):
-    if os.path.isdir(string) or re.match(r"(^\/|^\.\/|^\.\.\/|^[^/])[^:*?\"<>|\r\n]*\.yac$", string):
+    logger.debug("{}: string = {}".format("dir_path", string))
+    if os.path.isdir(string) or re.match(r"(^\/|^\.\/|^\.\.\/|^[^/])[^:*?\"<>|\r\n]*\.(yac|h)$", string):
         return string
     else:
-        raise Error(string)
+        raise TypeError("no valid path")
 
 def walk(args):
     logger = logging.getLogger(__name__)
@@ -17,27 +53,48 @@ def walk(args):
     logger.debug("Files: {}".format(files))
     logger.info("Number of files found: {}".format(len(files)))
 
-    if args["output"].endswith(".yac"):
+    if args["output"].endswith(YACFormatBinary.EXT) or args["output"].endswith(YACFormatCHeader.EXT):
         yd = YaraDatabase()
         for file in files:
             if file.endswith(".json"):
                 logger.info("Compiling file {}".format(file))
                 yd.add_file(file)
-        yd.write_file(args["output"])
+        data = yd.compile(
+            args["store_identifier_entry"],
+            args["store_identifier_signature"],
+            args["store_index_map_entries"],
+            args["store_index_map_signatures"],
+            args["store_index_map_string_blocks"],
+            args["store_hash"])
+        args["format"].write(args["output"], data)
     else:
         for file in files:
             if file.endswith(".json"):
                 logger.info("Compiling file {}".format(file))
                 yd = YaraDatabase()
                 yd.add_file(file)
-                yd.write_file(os.path.join(args["output"], os.path.splitext(os.path.basename(file))[0] + ".yac"))
+                data = yd.compile(
+                    args["store_identifier_entry"],
+                    args["store_identifier_signature"],
+                    args["store_index_map_entries"],
+                    args["store_index_map_signatures"],
+                    args["store_index_map_string_blocks"],
+                    args["store_hash"])
+                args["format"].write(os.path.join(args["output"], os.path.splitext(os.path.basename(file))[0] + ".yac"), data)
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='Compile single or multiple yara files')
     parser.add_argument('-i', '--input-directory', nargs='?', default='.', type=dir_path, help='Input directory (default: %(default)s)')
     parser.add_argument('-o', '--output', nargs='?', default='.', type=dir_path, help='Output file or directory (default: %(default)s)')
     parser.add_argument('-f', '--input-file', nargs='?', default='.', type=dir_path, help='Input file (default: %(default)s)')
+    parser.add_argument('--format', nargs='?', default=YACFormatBinary.STR, action='store', choices=[YACFormatBinary.STR, YACFormatCHeader.STR], help='Output file format (default: %(default)s)')
     parser.add_argument('-v', '--verbose', action="count", default=0, help="Verbosity level")
+    parser.add_argument('--store-identifier-entry', action='store_true', help='Store identifier for entry elements (default: %(default)s)')
+    parser.add_argument('--store-identifier-signature', action='store_true', help='Store identifier for signature elements (default: %(default)s)')
+    parser.add_argument('--store-index-map-entries', action='store_true', help='Store index map for entries (default: %(default)s)')
+    parser.add_argument('--store-index-map-signatures', action='store_true', help='Store identifier for signature elements (default: %(default)s)')
+    parser.add_argument('--store-index-map-string-blocks', action='store_true', help='Store identifier for string blocks (default: %(default)s)')
+    parser.add_argument('--store-hash', action='store_true', help='Store hash for database (default: %(default)s)')
     args = parser.parse_args()
 
     if args.verbose == 0:
@@ -54,7 +111,14 @@ if __name__ == "__main__":
         "input_directory": args.input_directory,
         "output": args.output,
         "input_file": args.input_file,
-        "verbosity": args.verbose
+        "format": DICT_STR_YACFORMAT[args.format],
+        "verbosity": args.verbose,
+        "store_identifier_entry": args.store_identifier_entry,
+        "store_identifier_signature": args.store_identifier_signature,
+        "store_index_map_entries": args.store_index_map_entries,
+        "store_index_map_signatures": args.store_index_map_signatures,
+        "store_index_map_string_blocks": args.store_index_map_string_blocks,
+        "store_hash": args.store_hash
     }
 
     logger.debug("args = {}".format(args))