import sys
import random
import string
import time
import itertools
import operator
import base64
  
from Crypto.Cipher import AES
from Crypto.Hash import SHA256

def slow_foo():
    p = 181
    k = 2
    while k < p:
        if p % k == 0:
            return
        k += 1
        
def aes_enc(k, m):
  """ 
  Encrypt a message m with a key k in ECB mode using AES as follows:
  c = AES(k, m)
  
  Args:
    m should be a bytestring multiple of 16 bytes (i.e. a sequence of characters such as 'Hello...' or '\x02\x04...')
    k should be a bytestring of length exactly 16 bytes.

  Return:
    The bytestring ciphertext c
  """
  aes = AES.new(k)
  c = aes.encrypt(m)

  return c

def aes_dec(k, c):
  """
  Decrypt a ciphertext c with a key k in ECB mode using AES as follows:
  m = AES(k, c)

  Args:
    c should be a bytestring multiple of 16 bytes (i.e. a sequence of characters such as 'Hello...' or '\x02\x04...')
    k should be a bytestring of length exactly 16 bytes.

  Return:
    The bytestring message m
  """
  aes = AES.new(k)
  m = aes.decrypt(c)
  
  return m

def aes_enc_cbc(k, m, iv):
  """
  Encrypt a message m with a key k in CBC mode using AES as follows:
  c = AES(k, m)

  Args:
    m should be a bytestring multiple of 16 bytes (i.e. a sequence of characters such as 'Hello...' or '\x02\x04...')
    k should be a bytestring of length exactly 16 bytes.
    iv should be a bytestring of length exactly 16 bytes.

  Return:
    The bytestring ciphertext c
  """
  aes = AES.new(k, AES.MODE_CBC, iv)
  c = aes.encrypt(m)

  return c

def aes_dec_cbc(k, c, iv):
  """
  Decrypt a ciphertext c with a key k in CBC mode using AES as follows:
  m = AES(k, c)

  Args:
    c should be a bytestring multiple of 16 bytes (i.e. a sequence of characters such as 'Hello...' or '\x02\x04...')
    k should be a bytestring of length exactly 16 bytes.
    iv should be a bytestring of length exactly 16 bytes.

  Return:
    The bytestring message m
  """
  aes = AES.new(k, AES.MODE_CBC, iv)
  m = aes.decrypt(c)

  return m
  
def aes_cbc_mac(k, m):
  """
  Compute a CBC-MAC of message m with a key k using AES as follows:
  t = AES-CBC-MAC(k=(k1,k2), m),
  where k1 is used for the raw-CBC operation and k2 is used for the final
  encryption.

  k1 and k2 are derived from k as follows:
  [k1|k2] = SHA256(k | "CBC MAC keys")

  Note: the IV for CBC in this case will be 0.

  Args:
    m should be a bytestring multiple of 16 bytes (i.e. a sequence of characters such as 'Hello...' or '\x02\x04...')
    k should be a bytestring of length exactly 16 bytes.

  Return:
    The bytestring MAC t, of 16 bytes.
  """

  #Require good size
  m = m.ljust(16)
  k = k.ljust(16)

  #Derive the keys for raw-CBC and for the final tag
  #[k1 | k2] = SHA256(k + "CBC MAC keys")

  #Get the MAC:
  #1 - Do aes-CBC with k1 and iv=0, then keep only last block (last 16 bytes) of encryption
  #2 - Perform another AES encryption (simple, without CBC) on the last block from #1 using k2
  #t = tag
  t = 16*'\x00'

  return t

def verify(message, tag):
  key = 'Cozonace si oua '

  # Get correct tag
  goodtag = aes_cbc_mac(key, message)
  
  # Compare tags
  for i in range(16):
    # Artificially extend byte comparison duration
    slow_foo()
    if tag[i] != goodtag[i]:
      return False

  return True


def main():
  message = 'Hristos a inviat'

  # Step 1. Iterate through all possible first byte values, and call the
  # Verify oracle for each of them
  tag = 16*'\x00'
  verify(message, tag)

  # Step 2. Store the byte that caused the longest computation time

  # Step 3. Continue the operation for each byte (except the last)

  # Step 4. Guess the last byte, and query the oracle with the complete tag
  mytag = '???'
  result = verify(message, mytag)
  if result == True:
    print "Found tag: " + mytag

if __name__ == "__main__":
  main()