require_relative '../util'

WEAK_HASH_IV = [0xf0, 0x0d].freeze
STRONGER_HASH_IV = [0x00, 0xc0, 0xff, 0xee].freeze

def weak_hash(buffer, iv = nil)
  aes_hash(buffer, iv || WEAK_HASH_IV)
end

def stronger_hash(buffer, iv = nil)
  aes_hash(buffer, iv || STRONGER_HASH_IV)
end

def concat_hash(buffer, iv = nil)
  weak_hash(buffer, iv) + stronger_hash(buffer, iv)
end

def find_collision_pair(iv)
  seen = {}
  loop do
    (0..255).to_a.repeated_combination(16) do |buffer|
      hash = yield(buffer, iv)
      return [buffer, seen[hash]] if seen.include?(hash)
      seen[hash] = buffer
    end
  end
end

def find_collisions(n, iv)
  pairs = []
  n.times do
    pair = find_collision_pair(iv) { |buffer, iv| yield(buffer, iv) }
    pairs << pair
    iv = yield(pair[0], iv)
  end
  pairs[0].product(*pairs.drop(1)).map { |pair| pair.flatten(1) }
end

def find_multicollision(collisions)
  seen = {}
  iters = 0
  collisions.each do |buffer|
    hash = yield(buffer)
    return [iters, buffer, seen[hash]] if seen.include?(hash)
    seen[hash] = buffer
    iters += 1
  end
  raise "didn't find multicollision"
end

bits = (STRONGER_HASH_IV.size / 2) * 8
collisions = find_collisions(bits, WEAK_HASH_IV) do |buffer, iv|
  aes_ecb_compress(buffer, iv)
end

info("generated #{collisions.size} collisions")
iters, m1, m2 = find_multicollision(collisions) { |buf| stronger_hash(buf) }
info("found multicollision after #{iters} iterations")
assert(weak_hash(m1) == weak_hash(m2))
assert(stronger_hash(m1) == stronger_hash(m2))
assert(concat_hash(m1) == concat_hash(m2))