#!/usr/bin/env python
import argparse
import json
import logging
import re
import traceback
import warnings

from regexploit.ast.sre import SreOpParser
from regexploit.bin.files import file_generator
from regexploit.output.text import TextOutput
from regexploit.redos import find


def get_json(filename: str):
    with open(filename, "rb") as f:
        try:
            return json.load(f)
        except json.decoder.JSONDecodeError:
            print(f"Error parsing JSON from {filename}")
            return


def handle_file(yamljson, filename: str, output: TextOutput):
    if isinstance(yamljson, (list, dict)):
        YamlJsonWalker(filename, output).handle(yamljson)


class YamlJsonWalker:
    def __init__(self, filename: str, output: TextOutput):
        self.filename = filename
        self.output = output

    def handle(self, elem):
        if isinstance(elem, str) and len(elem) > 5:
            try:
                parsed = SreOpParser().parse_sre(elem)
            except re.error:
                return  # We will have many strings which aren't actually regexes
            try:
                self.output.next()
                for redos in find(parsed):
                    if redos.starriness > 2:
                        self.output.record(
                            redos,
                            elem,
                            filename=self.filename,
                        )
            except Exception:
                print(f"Error finding ReDoS: {elem} from {self.filename}")
                print(traceback.format_exc())
        elif isinstance(elem, list):
            for _elem in elem:
                self.handle(_elem)
        elif isinstance(elem, dict):
            for _elem in elem.values():
                self.handle(_elem)


def main(get_object=get_json):
    with warnings.catch_warnings():
        # Some weird regexes emit warnings
        warnings.simplefilter("ignore", category=FutureWarning)
        warnings.simplefilter("ignore", category=DeprecationWarning)
        parser = argparse.ArgumentParser(
            description="Parse regexes out of YAML files (strings, lists and dictionary values) and scan them for ReDoS"
        )
        parser.add_argument("files", nargs="+", help="YAML files")
        parser.add_argument(
            "--glob", action="store_true", help="Glob the input filenames (**/*)"
        )
        parser.add_argument("--verbose", action="store_true", help="Verbose logging")
        parser.add_argument(
            "--ignore", action="append", help="Paths containing this string are ignored"
        )
        args = parser.parse_args()

        if args.verbose:
            logging.basicConfig(level=logging.DEBUG)

        files = file_generator(
            args.files,
            args.glob,
            ["*.json"] if get_object is get_json else ["*.yaml", "*.yml", "*.json"],
            args.ignore,
        )
        output = TextOutput()
        for filename in files:
            logging.debug(filename)
            handle_file(get_object(filename), filename, output)
        print(f"Processed {output.regexes} regexes")


def main_yaml():
    try:
        from yaml import safe_load, YAMLError

        def get_yaml(filename: str):
            with open(filename, "rb") as f:
                try:
                    return safe_load(f.read())
                except YAMLError:
                    print(f"Error parsing YAML from {filename}")
                    return

        main(get_object=get_yaml)
    except ImportError:
        print(
            "Pyyaml extra required: Install regexploit with 'pip install regexploit[yaml]' or run 'pip install pyyaml'"
        )
        raise


if __name__ == "__main__":
    main()
