import argparse
import os
import re
from yara import *

def dir_path(string):
    if os.path.isdir(string) or re.match(r"(^\/|^\.\/|^\.\.\/|^[^/])[^:*?\"<>|\r\n]*\.yac$", string):
        return string
    else:
        raise Error(string)

def walk(args):
    logger = logging.getLogger(__name__)
    logger.info("Walking files ...")

    files = [os.path.abspath(os.path.join(dp, f)) for dp, dn, filenames in os.walk(args["input_directory"]) for f in filenames]
    logger.debug("Files: {}".format(files))
    logger.info("Number of files found: {}".format(len(files)))

    if args["output"].endswith(".yac"):
        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"])
    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"))

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('-v', '--verbose', action="count", default=0, help="Verbosity level")
    args = parser.parse_args()

    if args.verbose == 0:
        log_level = logging.WARNING
    elif args.verbose == 1:
        log_level = logging.INFO
    elif args.verbose >= 2:
        log_level = logging.DEBUG

    logging.basicConfig(stream=sys.stdout, level=log_level)
    logger = logging.getLogger(__name__)

    args = {
        "input_directory": args.input_directory,
        "output": args.output,
        "input_file": args.input_file,
        "verbosity": args.verbose
    }

    logger.debug("args = {}".format(args))

    walk(args)