#!/usr/bin/env python
import argparse
import fileinput
import logging
import re
import sys
import traceback

from regexploit.ast.sre import SreOpParser
from regexploit.languages.javascript import fix_js_regex
from regexploit.output.text import TextOutput
from regexploit.redos import find


def find_redos(pattern: str, flags: int, output: TextOutput, parser):
    try:
        parsed = parser(pattern, flags)
    except Exception as e:
        print(f"Error parsing: {pattern}", e)
        return
    output.next()
    for redos in find(parsed):
        if redos.starriness > 2:
            output.record(redos, pattern)
            yield redos


def python(pattern: str, flags: int):
    return SreOpParser().parse_sre(pattern, flags)


def javascript(pattern: str, flags: int):
    try:
        return SreOpParser().parse_sre(pattern)
    except:
        try:
            fixed = fix_js_regex(pattern)
            re.compile(fixed)
        except:
            raise

        try:
            return SreOpParser().parse_sre(fixed)
        except:
            print(traceback.format_exc())
            raise


def main():
    parser = argparse.ArgumentParser(
        description="Parse regexes from stdin and scan them for ReDoS"
    )
    parser.add_argument(
        "-f",
        "--flavour",
        "--flavor",
        choices=["python", "js"],
        default="python",
        help="Regex language",
    )
    parser.add_argument(
        "-v", "--verbose", action="count", default=0, help="Verbose logging"
    )
    parser.add_argument(
        "-u",
        "--unescape",
        action="store_true",
        help="Unescape the regular expressions before parsing them (e.g. double backslashes)",
    )
    args = parser.parse_args()
    sys.argv = sys.argv[:1]
    if args.verbose == 1:
        logging.basicConfig(level=logging.INFO)
    elif args.verbose > 1:
        logging.basicConfig(level=logging.DEBUG)

    isatty = sys.stdin.isatty()
    if isatty:
        print("Welcome to Regexploit. Enter your regexes:")
    output = TextOutput()
    try:
        for line in fileinput.input():
            found = False
            line = line.rstrip("\n")
            if args.unescape:
                # \\d -> \d
                line = line.encode().decode("unicode_escape")
            for _ in find_redos(
                line, 0, output, javascript if args.flavour == "js" else python
            ):
                found = True
            if isatty and not found:
                print("No ReDoS found.")
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()
