from exploits.hashes.collisions import chicken_scheme
from test.exploits.dummy_output import DummyOutput


def truncate_64bit(x):
    return x & 0xFFFFFFFFFFFFFFFF


def chicken_scheme_hash(bytes, hash_table_size, initial_key):
    key = initial_key
    for byte in bytes:
        key ^= truncate_64bit(truncate_64bit(key << 6) + (key >> 2) + ord(byte))
    return key % hash_table_size


def test_run():
    output = DummyOutput()
    n_collisions = 10
    hash_table_size = 0x3FFFFFFFFFFFFFFF
    target = '42'
    initial_key = 0
    chicken_scheme.options['n_collisions'] = n_collisions
    chicken_scheme.options['n_substrings'] = 1
    chicken_scheme.options['initial_key'] = initial_key
    chicken_scheme.options['target_type'] = 'image'
    chicken_scheme.options['target'] = target
    chicken_scheme.options['hash_table_size'] = hash_table_size
    chicken_scheme.run(output)
    assert output.count() == n_collisions
    for i in output:
        assert chicken_scheme_hash(i, hash_table_size, initial_key) == int(target)


def test_more_collisions():
    output = DummyOutput()
    n_collisions = 100
    hash_table_size = 0x3FFFFFFFFFFFFFFF
    target = '42'
    initial_key = 0
    chicken_scheme.options['n_collisions'] = n_collisions
    chicken_scheme.options['n_substrings'] = 1
    chicken_scheme.options['initial_key'] = initial_key
    chicken_scheme.options['target_type'] = 'image'
    chicken_scheme.options['target'] = target
    chicken_scheme.options['hash_table_size'] = hash_table_size
    chicken_scheme.run(output)
    assert output.count() == n_collisions
    for i in output:
        assert chicken_scheme_hash(i, hash_table_size, initial_key) == int(target)


def test_preimage():
    output = DummyOutput()
    n_collisions = 10
    hash_table_size = 0x3FFFFFFFFFFFFFFF
    target = 'hello'
    initial_key = 0
    chicken_scheme.options['n_collisions'] = n_collisions
    chicken_scheme.options['n_substrings'] = 1
    chicken_scheme.options['initial_key'] = initial_key
    chicken_scheme.options['target_type'] = 'preimage'
    chicken_scheme.options['target'] = target
    chicken_scheme.options['hash_table_size'] = hash_table_size
    chicken_scheme.run(output)
    assert output.count() == n_collisions
    for i in output:
        assert (chicken_scheme_hash(i, hash_table_size, initial_key) ==
                chicken_scheme_hash(target, hash_table_size, initial_key))


def test_no_zero_intial_key():
    output = DummyOutput()
    n_collisions = 1
    hash_table_size = 0x3FFFFFFFFFFFFFFF
    target = '42'
    initial_key = 11312
    chicken_scheme.options['n_collisions'] = n_collisions
    chicken_scheme.options['n_substrings'] = 1
    chicken_scheme.options['initial_key'] = initial_key
    chicken_scheme.options['target_type'] = 'image'
    chicken_scheme.options['target'] = target
    chicken_scheme.options['hash_table_size'] = hash_table_size
    chicken_scheme.run(output)
    assert output.count() == n_collisions
    for i in output:
        assert chicken_scheme_hash(i, hash_table_size, initial_key) == int(target)
