import sys, os, itertools, json

from playcrypt.tools import *
from playcrypt.ideal.block_cipher import *
from playcrypt.ideal.message_authentication_code import *
from playcrypt.games.game_ufcma import GameUFCMA
from playcrypt.simulator.ufcma_sim import UFCMASim
from playcrypt.games.game_lr import GameLR
from playcrypt.simulator.lr_sim import LRSim
from playcrypt.games.game_int_ctxt import GameINTCTXT
from playcrypt.simulator.ctxt_sim import CTXTSim
from playcrypt.ideal.function_family import *


"""
Problem 2

Let E : {0,1}^k x {0,1}^{2n} -> {0,1}^{2n} be a block cipher with k, n >= 128.
Let K be the key-generation algorithm that returns a random k-bit key.
Let SE = (K,Enc,Dec) be the symmetric encryption scheme with encryption and
decryption algorithms as described below. Note that the message input to Enc
is an n-bit string, and the ciphertext input to Dec is a 4n-bit string.
"""

def Enc(K,M):
    if len(M) != n_bytes:
        return None

    A1 = random_string(n_bytes)
    A2 = xor_strings(M, A1)
    C = []
    C.append(E(K, (A1 + "\x00" * n_bytes)))
    C.append(E(K, (A2 + "\xFF" * n_bytes)))
    return join(C)

def Dec(K,C):
    if len(C) != 4 * n_bytes:
        return None

    C = split(C, 2 * n_bytes)
    X1 = E_I(K, C[0])            # X1 = A1 || P1 in the pseudocode
    X2 = E_I(K, C[1])            # X2 = A2 || P2 in the pseudocode
    X1 = split(X1, n_bytes)      # A1 is X1[0] ; P1 is X1[1]
    X2 = split(X2, n_bytes)      # A2 is X2[0] ; P2 is X2[1]

    if (X1[1] != "\x00" * n_bytes) or (X2[1] != "\xFF" * n_bytes):
        return None

    M = xor_strings(X1[0], X2[0])
    return M


"""
[10 points] Show that SE is not INT-CTXT secure by presenting an O(n)-time
adversary A2 making two queries with advantage at least 1 - 2^{-n}.
"""

def A2(enc):
    """You must fill in this method. We will define variables k, n, k_bytes,
    n_bytes, Enc and Dec for you.

    :param enc: This is the oracle supplied by the game.
    return: a forged ciphertext
    """
    M0 = int_to_string(0, n_bytes) # can be a random string
    R1 = random_string(n_bytes)
    c_t0 = enc(M0)
    c_t1 = enc(R1)
# NOTE  
    # Since 0^n and 1^n are repeating patterns for the encryption scheme and is validated from the decryption scheme
    # we can switch the encryption of the first ciphertext with the 2nd ciphertext, and it would still be valid
# NOTE
    c_t0 = split(c_t0, 2*n_bytes) 
    c_t1 = split(c_t1, 2*n_bytes)
    c_t0[1] = c_t1[1]

    return join(c_t0)



"""
==============================================================================================
The following lines are used to test your code, and should not be modified.
==============================================================================================
"""

if __name__ == '__main__':
    k = 128
    n = 128
    k_bytes = k//8
    n_bytes = n//8

    EE = BlockCipher(k_bytes, 2*n_bytes)
    E = EE.encrypt
    E_I = EE.decrypt

    g = GameINTCTXT(2, Enc, Dec, k_bytes)
    s = CTXTSim(g, A2)

    print ("When k=128, n=128:")
    print ("The advantage of your adversary A2 is ~" + str(s.compute_advantage()))

    k = 256
    n = 128
    k_bytes = k//8
    n_bytes = n//8

    EE = BlockCipher(k_bytes, 2*n_bytes)
    E = EE.encrypt
    E_I = EE.decrypt

    g = GameINTCTXT(2, Enc, Dec, k_bytes)
    s = CTXTSim(g, A2)

    print ("When k=256, n=128:")
    print ("The advantage of your adversary A2 is ~" + str(s.compute_advantage()))