#!/usr/bin/python3

# "Crypto" might need "pip install pycrypto" if it's not installed

import Crypto, random
from Crypto.Cipher import AES

def int_to_bytes(i,len): # Not optimized
    res = []
    for j in range(len):
        res.append(i%256)
        i = i>>8
    return bytes(res)

def aes_ecb_enc(k,m):
    from Crypto import Random
    assert isinstance(m,bytes)
    assert len(m)%AES.block_size == 0, len(m)%AES.block_size
    k = int_to_bytes(k,AES.block_size)
    cipher = AES.new(k, AES.MODE_ECB)
    return cipher.encrypt(m)

def aes_ecb_dec(k,m):
    from Crypto import Random
    assert isinstance(m,bytes)
    k = int_to_bytes(k,AES.block_size)
    cipher = AES.new(k, AES.MODE_ECB)
    return cipher.decrypt(m)

# Just a test
assert aes_ecb_dec(2123414234,aes_ecb_enc(2123414234,b'hello there test')) == b'hello there test'


# The game: it gets a prg and an adversary as arguments, 
# as well as the messages to be distinguished
def guessing_game(adv,m0,m1):
    b = random.randint(0,1) # Random bit
    k = random.getrandbits(256) # Random AES key
    seed = random.randint(0,2**32-1) # Random seed
    rand = [random.randint(0,2**32-1) for i in range(10)] # Truly random output
    msg = (m0,m1)[b]
    ciph = aes_ecb_enc(k,msg)
    badv = adv(ciph)
    return b==badv


def adv(ciph):
    # blocks contains the ciphertext as a list of blocks
    blocks = [ciph[i*AES.block_size:(i+1)*AES.block_size] for i in range(len(ciph)//AES.block_size)]

    ???
    return ??? # return 0 or 1

    
def test_adv(adv):
    num_true = 0
    num_tries = 3000
    m0 = open("ecb-distinguish-1.txt","rb").read()
    m1 = open("ecb-distinguish-2.txt","rb").read()
    for i in range(num_tries):
        #if i%100==0: print(str(i)+"...")
        if guessing_game(adv,m0,m1): num_true += 1
    ratio = float(num_true)/num_tries
    print(ratio)

# An output near 0.5 means no attack
# An output neat 0.0 or 1.0 means a successful attack
test_adv(adv)