expandRoundKeys new

def expandRoundKeys(key):
    """
    https://en.m.wikipedia.org/wiki/AES_key_schedule?fbclid=IwAR2MKdPC8cO1BeqKlkaheZYKKlUf_QgO7-t1SDPD5etsE8_RvhrLnQnWYAA
    """

    genRoundConsts()

    """
    N as the length of the key in 32-bit words:
    4 words for AES-128, 6 words for AES-192,
    and 8 words for AES-256

    R as the number of rounds:
    10 rounds for AES-128, 12 rounds for AES-192,
    14 rounds for AES-256
    """

    key_size = len(key)
    # print('key_size', key_size)
    num_keys = util.getNumOfRounds() + 1  # this is R

    round_keys = []

    N = key_size // 4
    for i in range(4 * num_keys):

        if i < N:
            # Initialize round keys with the original key
            currWord = []
            for j in range(4):
                currWord.append(key[i * 4 + j])
            round_keys.append(currWord)

            print('key', key)
            print('round keys', round_keys)

        elif i >= N and i % N == 0:
            prevWord = round_keys[i - 1].copy()
            prevWord = rot_word(prevWord)
            prevWord = sub_word(prevWord)
            prevWord[0] ^= roundConstsHex[i // N]
            # round constants are only added to the first element of a word
            prevWord = [round_keys[i - N][j] ^ prevWord[j] for j in range(4)]
            round_keys.append(prevWord)

        elif i >= N and N > 6 and i % N == 4:
            prevWord = round_keys[i - 1].copy()
            prevWord = sub_word(prevWord)  # SubWord(W[i-1])
            prevWord = [round_keys[i - N][j] ^ prevWord[j] for j in range(4)]
            round_keys.append(prevWord)

        else:
            prevWord = round_keys[i - 1].copy()
            prevWord = [round_keys[i - N][j] ^ prevWord[j] for j in range(4)]
            round_keys.append(prevWord)

        # print('round keys', round_keys)

    allRoundKeys = []
    for i in range(0, len(round_keys), 4):  # range(start, stop, step) 4 words per round
        currRoundKey = []
        for j in range(4):
            currRoundKey.append(round_keys[i + j])
        allRoundKeys.append(np.transpose(currRoundKey))

    flattened_list = [element for sublist in allRoundKeys for element in sublist]
    # print(flattened_list[2])
    global expandedRoundKeys
    expandedRoundKeys = flattened_list