require_relative '../util'
require 'zlib'

SESSION_ID = b64encode(random_bytes(32))
info("session ID: #{SESSION_ID}")

def format_request(input)
  "POST / HTTP/1.1
Host: hapless.com
Cookie: sessionid=#{SESSION_ID}
Content-Length: #{input.length}
#{input}
"
end

def compress(string)
  Zlib::Deflate.deflate(string)
end

def ctr_compression_oracle(string)
  key = random_bytes(16)
  nonce = rand(0...2**64)
  aes_ctr_encrypt(compress(format_request(string)).bytes, key, nonce).size
end

$last_length = nil

def report_progress(string)
  print " #{' ' * $last_length}\r" if $last_length
  print " #{string}\r"
  $last_length = string.length
end

CHARSET = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789+/='

def ctr_guess_byte(known)
  guesses = {}
  # this improves our chances considerably
  suffix = random_bytes(10, (128..255))
  CHARSET.each_byte do |byte|
    input = "sessionid=#{str(known + [byte] + suffix)}"
    guesses[byte] = ctr_compression_oracle(input)
  end
  guesses.minmax_by { |_, v| v }
end

def ctr_guess_bytes
  known = []
  loop do
    min, max = ctr_guess_byte(known)
    return known unless min[1] < max[1]
    known << min[0]
    report_progress(str(known))
  end
end

# borrowed from #32

MIN_ATTEMPTS = 2
MAX_ATTEMPTS = 25

def guess_repeatedly
  guesses = Hash.new { |h, k| h[k] = 0 }
  MIN_ATTEMPTS.times { guesses[yield] += 1 }
  attempts = MIN_ATTEMPTS
  loop do
    raise "couldn't guess solution" if attempts > MAX_ATTEMPTS
    best_guess = guesses.max_by { |_, v| v }
    return best_guess[0] if (best_guess[1] / attempts.to_f) > 0.5
    guesses[yield] += 1
    attempts += 1
  end
end

ctr_guess = str(guess_repeatedly { ctr_guess_bytes })
puts
assert(SESSION_ID == ctr_guess)

def cbc_compression_oracle(string)
  key = random_bytes(16)
  iv = random_bytes(16)
  buffer = pkcs7pad(compress(format_request(string)).bytes, 16)
  aes_cbc_encrypt(buffer, key, iv).size
end

def cbc_guess_block_size
  input_size = 0
  last_size = cbc_compression_oracle(random_bytes(input_size, (128..255)))
  loop do
    size = cbc_compression_oracle(random_bytes(input_size, (128..255)))
    return size - last_size if size > last_size
    input_size += 1
    last_size = size
  end
end

# you could do a lot better, but this is the safest approach...
def cbc_guess_byte(known, block_size)
  (0...block_size).each do |glue_size|
    suffix = random_bytes(glue_size, (128..255))
    guesses = {}
    CHARSET.each_byte do |byte|
      input = "sessionid=#{str(known + [byte] + suffix)}"
      guesses[byte] = cbc_compression_oracle(input)
    end
    min, max = guesses.minmax_by { |_, v| v }
    return min[0] if min[1] < max[1]
  end
  nil
end

def cbc_guess_bytes
  known = []
  block_size = cbc_guess_block_size
  loop do
    byte = cbc_guess_byte(known, block_size)
    return known unless byte
    known << byte
    report_progress(str(known))
  end
end

cbc_guess = str(cbc_guess_bytes)
puts
assert(SESSION_ID == cbc_guess)