#!/usr/bin/env ruby

require 'pwn'

$s = Sock.new('crypto.chal.csaw.io', 1003)
$s.recvuntil(':').split("\n")[0]
B = 16

def hexdecode(string)
  [string].pack('H*').bytes
end

def get_ciphertext(input)
  puts "Trying #{input}"
  $s.sendline(input)
  $s.recvline
  hexdecode($s.recvline.chomp)
end

def duplicate_blocks(ciphertext)
  blocks = ciphertext.each_slice(B)
  blocks.each_with_index do |block, i|
    return i if blocks.select { |b| b == block }.length > 1
  end
  false
end

def guess_prefix_size
  len = 0
  glue = 'X' * 32
  loop do
    i = duplicate_blocks(get_ciphertext('A' * len + glue))
    return B * i - len if i
    len += 1
  end
end

def guess_suffix_size(prefix_size)
  input = 'X' * 16
  old_len = get_ciphertext(input).length
  loop do
    input += 'X'
    new_len = get_ciphertext(input).length
    return new_len - B - prefix_size - input.length if new_len > old_len
  end
end

prefix_size = guess_prefix_size
glue_size = B - (prefix_size % B)
suffix_size = guess_suffix_size(prefix_size)
puts "Guessed prefix size: #{prefix_size}"
puts "Guessed glue size: #{glue_size}"
puts "Guessed suffix size: #{suffix_size}"

def make_short_blocks(prefix_size, glue_size, known)
  block_count = ((known.length + 1) / B.to_f).ceil
  block = 'A' * (block_count * B - known.length - 1)
  glue = 'X' * glue_size
  ciphertext = get_ciphertext(glue + block)
  ciphertext.drop(prefix_size + glue_size).take(B * block_count)
end

def find_in_dict(prefix_size, glue_size, known, short_blocks)
  block_count = ((known.length + 1) / B.to_f).ceil
  block = 'A' * (block_count * B - known.length - 1)
  glue = 'X' * glue_size
  (32..126).reverse_each do |b|
    ciphertext = get_ciphertext(glue + block + known + b.chr)
    match = ciphertext.drop(prefix_size + glue_size).take(B * block_count)
    padding = ciphertext[-B..-1]
    p [b.chr, padding]
    return b.chr if match == short_blocks
  end
  raise ':<'
end

known = ''
suffix_size.times do
  blocks = make_short_blocks(prefix_size, glue_size, known)
  puts "Created blocks"
  known += find_in_dict(prefix_size, glue_size, known, blocks)
  puts "Guessed: #{known}"
end

$s.close