require_relative '../util'

class Server
  def initialize(key)
    @key = key
  end

  def process(message)
    return 'message too short' unless message.length > 64
    mac = message.slice(-32, 32)
    iv = hexdecode(message.slice(-64, 32))
    buffer = message[0...-64].bytes
    params = decode_query_string(message[0...-64])
    return 'invalid message' unless params.keys == %w(from to amount)
    from, to, amount = params.values
    actual_mac = cbc_mac(pkcs7pad(buffer, 16), @key, iv)
    return "transacting #{amount}SB from #{from} to #{to}" if mac == actual_mac
    'invalid mac, transaction aborted'
  end
end

class Client
  def initialize(key)
    @key = key
  end

  def make_message(params)
    message = encode_query_string(params)
    iv = random_bytes(16)
    mac = cbc_mac(pkcs7pad(message.bytes, 16), @key, iv)
    message + hexencode(iv) + mac
  end
end

def make_server_and_client
  key = random_bytes(16)
  [Server.new(key), Client.new(key)]
end

# we'll want to mutate
#      from=anne&to=bob|&amount=1000000
# into
#      from=anne&to=eve|&amount=1000000
# while still keeping the MAC intact

def mess_with_iv
  server, client = make_server_and_client
  message = client.make_message(from: 'anne', to: 'bob', amount: 1_000_000)
  mac = message.slice(-32, 32)
  iv = hexdecode(message.slice(-64, 32))
  buffer = message[0...-64].bytes
  [[13, 14, 15], 'bob'.bytes, 'eve'.bytes].transpose.each do |i, a, b|
    buffer[i] ^= a ^ b
    iv[i] ^= a ^ b
  end
  new_message = str(buffer) + hexencode(iv) + mac
  server.process(new_message)
end

status = mess_with_iv
info(status)
assert(status[/to eve/])

class BetterServer
  def initialize(key)
    @key = key
    @iv = Array.new(16, 0)
  end

  def process(message)
    return 'message too short' unless message.length > 32
    mac = message.slice(-32, 32)
    buffer = message[0...-32].bytes
    info(message[0...-32].inspect)
    params = decode_query_string(message[0...-32])
    return 'invalid message' unless params.keys == %w(from tx_list)
    actual_mac = cbc_mac(pkcs7pad(buffer, 16), @key, @iv)
    from, tx_list = params.values
    transactions = tx_list.split(';').map { |kv| kv.split(':') }
    if mac == actual_mac
      transactions.map do |to, amount|
        "transacting #{amount}SB from #{from} to #{to}"
      end.join("\n")
    else
      'invalid mac, transaction aborted'
    end
  end
end

class BetterClient
  def initialize(key)
    @key = key
    @iv = Array.new(16, 0)
  end

  def make_message(params)
    message = encode_query_string(params)
    mac = cbc_mac(pkcs7pad(message.bytes, 16), @key, @iv)
    message + mac
  end
end

def make_better_server_and_client
  key = random_bytes(16)
  [BetterServer.new(key), BetterClient.new(key)]
end

def extension_attack
  server, client = make_better_server_and_client
  # captured from anne
  params1 = { from: 'anne', tx_list: 'bob:10;bob:11' }
  m1 = encode_query_string(params1).bytes
  t1 = hexdecode(client.make_message(params1).slice(-32, 32))
  # captured from another eve-controlled account, e
  params2 = { from: 'e', tx_list: 'a:0;eve:1000000' }
  m2 = encode_query_string(params2).bytes
  t2 = hexdecode(client.make_message(params2).slice(-32, 32))
  t1.each_with_index { |b, i| m2[i] ^= b }
  message = str(pkcs7pad(m1, 16) + m2) + hexencode(t2)
  server.process(message)
end

status = extension_attack
status.split("\n").each { |line| info(line) }
assert(status[/to eve/])