"""
The util module contains various utility functions.
"""

from __future__ import print_function
import argparse
from six.moves import range, cPickle
import re
import binascii
import sys
import pwnypack.main
import pwnypack.codec


__all__ = [
    'cycle',
    'cycle_find',
    'reghex',
    'pickle_call',
]


def deBruijn(n, k):
    """
    An implementation of the FKM algorithm for generating the de Bruijn
    sequence containing all k-ary strings of length n, as described in
    "Combinatorial Generation" by Frank Ruskey.
    """

    a = [ 0 ] * (n + 1)

    def gen(t, p):
        if t > n:
            for v in a[1:p + 1]:
                yield v
        else:
            a[t] = a[t - p]
         
            for v in gen(t + 1, p):
                yield v
         
            for j in range(a[t - p] + 1, k):
                a[t] = j
                for v in gen(t + 1, t):
                    yield v

    return gen(1, 1)


def cycle(length, width=4):
    """
    Generate a de Bruijn sequence of a given length (and width). A de Bruijn
    sequence is a set of varying repetitions where each sequence of *n*
    characters is unique within the sequence. This type of sequence can be
    used to easily find the offset to the return pointer when exploiting a
    buffer overflow.

    Args:
        length(int): The length of the sequence to generate.
        width(int): The width of each element in the sequence.

    Returns:
        str: The sequence.

    Example:
        >>> from pwny import *
        >>> cycle(80)
        AAAABAAACAAADAAAEAAAFAAAGAAAHAAAIAAAJAAAKAAALAAAMAAANAAAOAAAPAAAQAAARAAASAAATAAA
    """

    iter = deBruijn(width, 26)
    return ''.join([chr(ord('A') + next(iter)) for i in range(length)])


def cycle_find(key, width=4):
    """
    Given an element of a de Bruijn sequence, find its index in that sequence.

    Args:
        key(str): The piece of the de Bruijn sequence to find.
        width(int): The width of each element in the sequence.

    Returns:
        int: The index of ``key`` in the de Bruijn sequence.
    """

    key_len = len(key)
    buf = ''

    it = deBruijn(width, 26)

    for i in range(key_len):
        buf += chr(ord('A') + next(it))

    if buf == key:
        return 0

    for i, c in enumerate(it):
        buf = buf[1:] + chr(ord('A') + c)
        if buf == key:
            return i + 1

    return -1


REGHEX_PATTERN = r'(([a-fA-F0-9]{2})|(([?.])(\{\d+\})?)|(\*|\+)|\s+)'
reghex_check = re.compile(REGHEX_PATTERN + '+')
reghex_regex = re.compile(REGHEX_PATTERN)


def reghex(pattern):
    """
    Compile a regular hexpression (a short form regular expression subset
    specifically designed for searching for binary strings).

    A regular hexpression consists of hex tuples interspaced with control
    characters. The available control characters are:

    - ``?``: Any byte (optional).
    - ``.``: Any byte (required).
    - ``?{n}``: A set of 0 up to *n* bytes.
    - ``.{n}``: A set of exactly *n* bytes.
    - ``*``: Any number of bytes (or no bytes at all).
    - ``+``: Any number of bytes (at least one byte).

    Args:
        pattern(str): The reghex pattern.

    Returns:
        regexp: A regular expression as returned by ``re.compile()``.
    """

    if not reghex_check.match(pattern):
        raise SyntaxError('Invalid reghex pattern.')

    b_pattern = b''

    for match in reghex_regex.finditer(pattern):
        _, match_hex, _, match_char, match_char_len, match_star_plus = match.groups()
        if match_hex:
            b_pattern += pwnypack.codec.dehex(match_hex)
        elif match_char:
            if match_char == '?':
                if match_char_len is None:
                    b_pattern += b'.?'
                else:
                    b_pattern += ('.{0,%d}?' % int(match_char_len[1:-1])).encode('ascii')
            else:
                if match_char_len is None:
                    b_pattern += b'.'
                else:
                    b_pattern += b'.' * int(match_char_len[1:-1])
        elif match_star_plus:
            b_pattern += b'.' + match_star_plus.encode('ascii') + b'?'

    try:
        return re.compile(b_pattern)
    except (TypeError, binascii.Error, re.error):
        raise SyntaxError('Invalid reghex pattern.')


def pickle_call(func, *args):
    """
    Create a byte sequence which when unpickled calls a callable with given
    arguments.

    Arguments:
        func(callable): The function to call or class to instantiate.
        args(tuple): The arguments to call the callable with.

    Returns:
        bytes: The data that when unpickled calls ``func(*args)``.

    Example:
        >>> from pwny import *
        >>> import pickle
        >>> def hello(arg):
        ...     print('Hello, %s!' % arg)
        ...
        >>> pickle.loads(pickle_call(hello, 'world'))
        Hello, world!
    """

    class CallFunc(object):
        def __reduce__(self):
            return func, tuple(args)
    return cPickle.dumps(CallFunc(), 0)


@pwnypack.main.register('cycle')
def cycle_app(parser, cmd, args):  # pragma: no cover
    """
    Generate a de Bruijn sequence of a given length.
    """

    parser.add_argument('-w', '--width', type=int, default=4, help='the length of the cycled value')
    parser.add_argument('length', type=int, help='the cycle length to generate')
    args = parser.parse_args(args)
    return cycle(args.length, args.width)


@pwnypack.main.register('cycle-find')
def cycle_find_app(_parser, cmd, args):  # pragma: no cover
    """
    Find the first position of a value in a de Bruijn sequence.
    """

    parser = argparse.ArgumentParser(
        prog=_parser.prog,
        description=_parser.description,
    )
    parser.add_argument('-w', '--width', type=int, default=4, help='the length of the cycled value')
    parser.add_argument('value', help='the value to determine the position of, read from stdin if missing', nargs='?')
    args = parser.parse_args(args)
    index = cycle_find(pwnypack.main.string_value_or_stdin(args.value), args.width)
    if index == -1:
        print('Not found.')
        sys.exit(1)
    else:
        print('Found at position: %d' % index)
