"""

I am going to try to write the AES Algorithm from scratch without additional libraries in Python

Encrypt Steps:
1. Accept inputs
    1a. Accept Plaintext
    1b. Accept Encryption Key
        add padding if key is not 128 bits (16 ASCII characters in 128 bits)
2. AddRoundKey     ## Is this the initial key
3. Loop (10x,except on last loop, no MixColumn)
    ---
    3a. ByteSub (10x)
    3b. ShiftRow (10x)
    3c. MixColumn (9x)
    3d. AddRoundKey (10x)
    ---
4. Result is the Ciphertext

Decrypt Steps:
Decode is to invert each step above. ## Not included in this code

"""


"""
Predefined AES specific matrices
"""

s_box = [
    0x63, 0x7C, 0x77, 0x7B, 0xF2, 0x6B, 0x6F, 0xC5, 0x30, 0x01, 0x67, 0x2B, 0xFE, 0xD7, 0xAB, 0x76,
    0xCA, 0x82, 0xC9, 0x7D, 0xFA, 0x59, 0x47, 0xF0, 0xAD, 0xD4, 0xA2, 0xAF, 0x9C, 0xA4, 0x72, 0xC0,
    0xB7, 0xFD, 0x93, 0x26, 0x36, 0x3F, 0xF7, 0xCC, 0x34, 0xA5, 0xE5, 0xF1, 0x71, 0xD8, 0x31, 0x15,
    0x04, 0xC7, 0x23, 0xC3, 0x18, 0x96, 0x05, 0x9A, 0x07, 0x12, 0x80, 0xE2, 0xEB, 0x27, 0xB2, 0x75,
    0x09, 0x83, 0x2C, 0x1A, 0x1B, 0x6E, 0x5A, 0xA0, 0x52, 0x3B, 0xD6, 0xB3, 0x29, 0xE3, 0x2F, 0x84,
    0x53, 0xD1, 0x00, 0xED, 0x20, 0xFC, 0xB1, 0x5B, 0x6A, 0xCB, 0xBE, 0x39, 0x4A, 0x4C, 0x58, 0xCF,
    0xD0, 0xEF, 0xAA, 0xFB, 0x43, 0x4D, 0x33, 0x85, 0x45, 0xF9, 0x02, 0x7F, 0x50, 0x3C, 0x9F, 0xA8,
    0x51, 0xA3, 0x40, 0x8F, 0x92, 0x9D, 0x38, 0xF5, 0xBC, 0xB6, 0xDA, 0x21, 0x10, 0xFF, 0xF3, 0xD2,
    0xCD, 0x0C, 0x13, 0xEC, 0x5F, 0x97, 0x44, 0x17, 0xC4, 0xA7, 0x7E, 0x3D, 0x64, 0x5D, 0x19, 0x73,
    0x60, 0x81, 0x4F, 0xDC, 0x22, 0x2A, 0x90, 0x88, 0x46, 0xEE, 0xB8, 0x14, 0xDE, 0x5E, 0x0B, 0xDB,
    0xE0, 0x32, 0x3A, 0x0A, 0x49, 0x06, 0x24, 0x5C, 0xC2, 0xD3, 0xAC, 0x62, 0x91, 0x95, 0xE4, 0x79,
    0xE7, 0xC8, 0x37, 0x6D, 0x8D, 0xD5, 0x4E, 0xA9, 0x6C, 0x56, 0xF4, 0xEA, 0x65, 0x7A, 0xAE, 0x08,
    0xBA, 0x78, 0x25, 0x2E, 0x1C, 0xA6, 0xB4, 0xC6, 0xE8, 0xDD, 0x74, 0x1F, 0x4B, 0xBD, 0x8B, 0x8A,
    0x70, 0x3E, 0xB5, 0x66, 0x48, 0x03, 0xF6, 0x0E, 0x61, 0x35, 0x57, 0xB9, 0x86, 0xC1, 0x1D, 0x9E,
    0xE1, 0xF8, 0x98, 0x11, 0x69, 0xD9, 0x8E, 0x94, 0x9B, 0x1E, 0x87, 0xE9, 0xCE, 0x55, 0x28, 0xDF,
    0x8C, 0xA1, 0x89, 0x0D, 0xBF, 0xE6, 0x42, 0x68, 0x41, 0x99, 0x2D, 0x0F, 0xB0, 0x54, 0xBB, 0x16
]

Mul_2 = [
    0x00, 0x02, 0x04, 0x06, 0x08, 0x0a, 0x0c, 0x0e, 0x10, 0x12, 0x14, 0x16, 0x18, 0x1a, 0x1c, 0x1e,
    0x20, 0x22, 0x24, 0x26, 0x28, 0x2a, 0x2c, 0x2e, 0x30, 0x32, 0x34, 0x36, 0x38, 0x3a, 0x3c, 0x3e,
    0x40, 0x42, 0x44, 0x46, 0x48, 0x4a, 0x4c, 0x4e, 0x50, 0x52, 0x54, 0x56, 0x58, 0x5a, 0x5c, 0x5e,
    0x60, 0x62, 0x64, 0x66, 0x68, 0x6a, 0x6c, 0x6e, 0x70, 0x72, 0x74, 0x76, 0x78, 0x7a, 0x7c, 0x7e,
    0x80, 0x82, 0x84, 0x86, 0x88, 0x8a, 0x8c, 0x8e, 0x90, 0x92, 0x94, 0x96, 0x98, 0x9a, 0x9c, 0x9e,
    0xa0, 0xa2, 0xa4, 0xa6, 0xa8, 0xaa, 0xac, 0xae, 0xb0, 0xb2, 0xb4, 0xb6, 0xb8, 0xba, 0xbc, 0xbe,
    0xc0, 0xc2, 0xc4, 0xc6, 0xc8, 0xca, 0xcc, 0xce, 0xd0, 0xd2, 0xd4, 0xd6, 0xd8, 0xda, 0xdc, 0xde,
    0xe0, 0xe2, 0xe4, 0xe6, 0xe8, 0xea, 0xec, 0xee, 0xf0, 0xf2, 0xf4, 0xf6, 0xf8, 0xfa, 0xfc, 0xfe,
    0x1b, 0x19, 0x1f, 0x1d, 0x13, 0x11, 0x17, 0x15, 0x0b, 0x09, 0x0f, 0x0d, 0x03, 0x01, 0x07, 0x05,
    0x3b, 0x39, 0x3f, 0x3d, 0x33, 0x31, 0x37, 0x35, 0x2b, 0x29, 0x2f, 0x2d, 0x23, 0x21, 0x27, 0x25,
    0x5b, 0x59, 0x5f, 0x5d, 0x53, 0x51, 0x57, 0x55, 0x4b, 0x49, 0x4f, 0x4d, 0x43, 0x41, 0x47, 0x45,
    0x7b, 0x79, 0x7f, 0x7d, 0x73, 0x71, 0x77, 0x75, 0x6b, 0x69, 0x6f, 0x6d, 0x63, 0x61, 0x67, 0x65,
    0x9b, 0x99, 0x9f, 0x9d, 0x93, 0x91, 0x97, 0x95, 0x8b, 0x89, 0x8f, 0x8d, 0x83, 0x81, 0x87, 0x85,
    0xbb, 0xb9, 0xbf, 0xbd, 0xb3, 0xb1, 0xb7, 0xb5, 0xab, 0xa9, 0xaf, 0xad, 0xa3, 0xa1, 0xa7, 0xa5,
    0xdb, 0xd9, 0xdf, 0xdd, 0xd3, 0xd1, 0xd7, 0xd5, 0xcb, 0xc9, 0xcf, 0xcd, 0xc3, 0xc1, 0xc7, 0xc5,
    0xfb, 0xf9, 0xff, 0xfd, 0xf3, 0xf1, 0xf7, 0xf5, 0xeb, 0xe9, 0xef, 0xed, 0xe3, 0xe1, 0xe7, 0xe5,
]

Mul_3 = [
    0x00, 0x03, 0x06, 0x05, 0x0c, 0x0f, 0x0a, 0x09, 0x18, 0x1b, 0x1e, 0x1d, 0x14, 0x17, 0x12, 0x11,
    0x30, 0x33, 0x36, 0x35, 0x3c, 0x3f, 0x3a, 0x39, 0x28, 0x2b, 0x2e, 0x2d, 0x24, 0x27, 0x22, 0x21,
    0x60, 0x63, 0x66, 0x65, 0x6c, 0x6f, 0x6a, 0x69, 0x78, 0x7b, 0x7e, 0x7d, 0x74, 0x77, 0x72, 0x71,
    0x50, 0x53, 0x56, 0x55, 0x5c, 0x5f, 0x5a, 0x59, 0x48, 0x4b, 0x4e, 0x4d, 0x44, 0x47, 0x42, 0x41,
    0xc0, 0xc3, 0xc6, 0xc5, 0xcc, 0xcf, 0xca, 0xc9, 0xd8, 0xdb, 0xde, 0xdd, 0xd4, 0xd7, 0xd2, 0xd1,
    0xf0, 0xf3, 0xf6, 0xf5, 0xfc, 0xff, 0xfa, 0xf9, 0xe8, 0xeb, 0xee, 0xed, 0xe4, 0xe7, 0xe2, 0xe1,
    0xa0, 0xa3, 0xa6, 0xa5, 0xac, 0xaf, 0xaa, 0xa9, 0xb8, 0xbb, 0xbe, 0xbd, 0xb4, 0xb7, 0xb2, 0xb1,
    0x90, 0x93, 0x96, 0x95, 0x9c, 0x9f, 0x9a, 0x99, 0x88, 0x8b, 0x8e, 0x8d, 0x84, 0x87, 0x82, 0x81,
    0x9b, 0x98, 0x9d, 0x9e, 0x97, 0x94, 0x91, 0x92, 0x83, 0x80, 0x85, 0x86, 0x8f, 0x8c, 0x89, 0x8a,
    0xab, 0xa8, 0xad, 0xae, 0xa7, 0xa4, 0xa1, 0xa2, 0xb3, 0xb0, 0xb5, 0xb6, 0xbf, 0xbc, 0xb9, 0xba,
    0xfb, 0xf8, 0xfd, 0xfe, 0xf7, 0xf4, 0xf1, 0xf2, 0xe3, 0xe0, 0xe5, 0xe6, 0xef, 0xec, 0xe9, 0xea,
    0xcb, 0xc8, 0xcd, 0xce, 0xc7, 0xc4, 0xc1, 0xc2, 0xd3, 0xd0, 0xd5, 0xd6, 0xdf, 0xdc, 0xd9, 0xda,
    0x5b, 0x58, 0x5d, 0x5e, 0x57, 0x54, 0x51, 0x52, 0x43, 0x40, 0x45, 0x46, 0x4f, 0x4c, 0x49, 0x4a,
    0x6b, 0x68, 0x6d, 0x6e, 0x67, 0x64, 0x61, 0x62, 0x73, 0x70, 0x75, 0x76, 0x7f, 0x7c, 0x79, 0x7a,
    0x3b, 0x38, 0x3d, 0x3e, 0x37, 0x34, 0x31, 0x32, 0x23, 0x20, 0x25, 0x26, 0x2f, 0x2c, 0x29, 0x2a,
    0x0b, 0x08, 0x0d, 0x0e, 0x07, 0x04, 0x01, 0x02, 0x13, 0x10, 0x15, 0x16, 0x1f, 0x1c, 0x19, 0x1a,
]

mix_matrix = [
    [2, 3, 1, 1],
    [1, 2, 3, 1],
    [1, 1, 2, 3],
    [3, 1, 1, 2]
]

rci = [0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1B, 0x36]


def add_padding(input):
    """ Comments:
    0123456 ==> 01234560000000000 and 1234567890AB ==> 1234567890AB0000
    padding formula in action:
    Creates new variable "padded"
    initialKey plus..
    Takes the remainder of... the length of the initialKey "ToBePadded" divided by 16
    and subtracts that from 16
    This number is multiplied by the string "0" and appended onto the end of the initialKey
    0 character is considered as padding in encryption (AES).
    End result is a key that is always divisible by 16 characters
    """
    if len(input) % 16 == 0:
        return input
    else:
        padded = input + (16 - len(input) % 16) * "0"
        # print(padded)
        return padded


def grids_16(s):
    """
    Initial input --> 0123456789ABCDEF
    Should be broken down into this 4 by 4 matrix
    0   4   8   C
    1   5   9   D
    2   6   A   E
    3   7   B   F
    """
    all_blocks = []

    # Determine how many blocks of 16 bytes are needed
    # Creates matrix in Column Major
    for i in range(len(s) // 16):
        b = s[i * 16: i * 16 + 16]
        grid = [[], [], [], []]
        for j in range(4):
            for k in range(4):
                grid[j].append(b[j + k * 4])
        all_blocks.append(grid)
    return all_blocks


def generate_round_keys(key):
    keys = key
    for i in range(4, 4 * 11):
        next_column = []
        if i % 4 == 0:
            column_rotated = rot_words(keys[i - 1])
            column_s_boxed = []
            for val in column_rotated:
                column_s_boxed.append(s_box[val])

            round_column = [0, 0, 0, 0]
            round_column[0] = rci[int(i / 4) - 1]

            for j in range(4):
                temp = keys[i - 4][j] ^ column_s_boxed[j] ^ round_column[j]
                next_column.append(temp)

        else:
            for j in range(4):
                temp = keys[i - 4][j] ^ keys[i - 1][j]
                next_column.append(temp)
        keys.append(next_column)
    return keys


def rot_words(column):
    temp_words = [[], [], [], []]
    temp_words[0] = column[1]
    temp_words[1] = column[2]
    temp_words[2] = column[3]
    temp_words[3] = column[0]

    return temp_words


def add_round_key(pt, rk):
    # print(' -----\n     Adding Round Key Steps:')
    for i in range(len(pt)):
        for j in range(4):
            for k in range(4):
                x = pt[i][j][k]
                y = rk[j][k]
                z = int(hex(x ^ y), 0)
                pt[i][j][k] = z
                # print(pt)
    return pt


def byte_sub(a_blocklist):
    """
    0   1   2   3
    4   5   6   7   ==> Replace with s_box values
    8   9   A   B
    C   D   E   F
    """
    # print(" -----\n     Byte Sub Steps:")
    b_block_list = []  # Creates a new blank matrix called B
    for block in a_blocklist:  # The first block is the first 128 bit segment
        b_matrix = []  # Creates a new blank matrix called b_matrix
        for row in block:  # For "Hello World" the first ROW is [H, o, r, 0]
            b_matrix.append([])  # Added a new blank list into the next cell of the higher list
            # This will hold the next ROW of a_blockList subbed with s_box
            for val in row:  # ['0x48', '0x6f', '0x72', '0x30'] first ROW values in a_blockList
                # Converts the input value, which is a string hex value, to an integer base 16
                # This integer finds the index of the flat-list s_box
                # Converts this back to a string hex value
                # appends this string hex value into the last list item in b_matrix
                b_matrix[-1].append(s_box[val])
        b_block_list.append(b_matrix)
        # print(b_block_list)

    return b_block_list


def shift_row(pt):
    """
    0   1   2   3           0   1   2   3
    4   5   6   7   ==>     5   6   7   4
    8   9   A   B           A   B   8   9
    C   D   E   F           F   C   D   E
    """
    # print(" -----\n     Shift Row Steps:")
    for i in range(len(pt)):
        pt[i][0][1], pt[i][1][1], pt[i][2][1], pt[i][3][1] = pt[i][1][1], pt[i][2][1], pt[i][3][1], pt[i][0][1]
        pt[i][0][2], pt[i][1][2], pt[i][2][2], pt[i][3][2] = pt[i][2][2], pt[i][3][2], pt[i][0][2], pt[i][1][2]
        pt[i][0][3], pt[i][1][3], pt[i][2][3], pt[i][3][3] = pt[i][3][3], pt[i][0][3], pt[i][1][3], pt[i][2][3]
        # print(pt)
    c_blocklist = pt
    return c_blocklist


def mix_column(pt):
    """
    2   3   1   1
    1   2   3   1
    1   1   2   3
    3   1   1   2

    """

    # print(" -----\n     Mix Column Steps:")
    d_blocklist = []
    for block in range(len(pt)):
        d_matrix = []
        for column in range(4):
            d_matrix_col = []
            for mix_row in range(4):
                d_matrix_cell_xor = []
                for cell in range(4):
                    if mix_matrix[mix_row][cell] == 1:
                        temp = pt[block][column][cell]
                    if mix_matrix[mix_row][cell] == 2:
                        # temp = pt[block][column][cell] % 256
                        temp = Mul_2[pt[block][column][cell]]
                    if mix_matrix[mix_row][cell] == 3:
                        # temp = ((2 * pt[block][column][cell]) ^ (pt[block][column][cell])) % 256
                        temp = Mul_3[pt[block][column][cell]]
                    d_matrix_cell_xor.append(temp)
                d_matrix_cell = d_matrix_cell_xor[0] ^ d_matrix_cell_xor[1] ^ d_matrix_cell_xor[2] ^ d_matrix_cell_xor[
                    3]
                d_matrix_col.append(d_matrix_cell)
            d_matrix.append(d_matrix_col)
            # print(d_matrix)
        d_blocklist.append(d_matrix)
    return d_blocklist


''' Running The Code '''
if __name__ == '__main__':

    plaintext = input("Input the message you would like to encrypt: ")
    initialKey = input("Input your encryption key/password (up to 128 bits): ")  # accepts an input key of any length

    '''
    Trying to figure out how to add padding to the plaintext
    and then only AES the first 128 bits.
    Then repeat for each 128 bit section of the plaintext
    Need to do some kind of looping on the whole program after I write it
    '''

    # PADDING
    pt_padded = add_padding(plaintext)  # Padding the plaintext if not in 128 bit increments
    print("   -----------\nThe Padded PLAINTEXT is: ")
    print(pt_padded)
    # print("pt_Padded = " + pt_Padded)

    key_padded = add_padding(initialKey)  # Performs padding function of adding 0s to end
    print("   -----------\nThe Padded Encryption Key is: ")  # Prints the padded key with 0s added to the end
    print(key_padded)

    # From String to Integers
    key_ints = [ord(i) for i in key_padded]
    pt_ints = [ord(i) for i in pt_padded]  # Turning the plaintext into corresponding int values

    # Into 4 by 4 grid of integers
    key_ints_grid = grids_16(key_ints)
    # print(key_ints_grid)
    key_128bit = key_ints_grid[0]  # breaks key down into simply 1 4 by 4 matrix
    print("   -----------\nYour initial KEY in a padded 4 by 4 (16 byte) int grid is:")
    print(key_128bit)

    pt_ints_grid = grids_16(pt_ints)  # Organizing those hex values into a 4 by 4 grid (16 bytes)
    print("   -----------\nYour initial PLAINTEXT in padded BLOCKS of 4 by 4 (16 byte) in grid is:")
    print(pt_ints_grid)

    after_round_key = add_round_key(pt_ints_grid, key_128bit)
    pt_grid = after_round_key
    print("   PLAINTEXT after Add Round Key:")
    print(after_round_key)


    roundKeys = generate_round_keys(key_128bit)

    for i in range(10):
        pt_byte_sub_grid = byte_sub(pt_grid)
        print("   PLAINTEXT after Byte Substitution:")
        print(pt_byte_sub_grid)
        pt_grid = pt_byte_sub_grid

        after_shift_row = shift_row(pt_grid)
        print("   PLAINTEXT after Shift Row:")
        print(after_shift_row)
        pt_grid = after_shift_row


        if i < 9:
            after_mix_column = mix_column(pt_grid)
            print("   PLAINTEXT after Mix Column:")
            print(after_mix_column)
            pt_grid = after_mix_column


        key_Block = roundKeys[4 * (i + 1): (4 * (i + 1)) + 4]  # Next 4 columns

        after_round_key = add_round_key(pt_grid, key_Block)
        print("   PLAINTEXT after Add Round Key:")
        print(after_round_key)
        pt_grid = after_round_key

        print("   ROUND" + str(i + 1) + " Complete\n*********************")
    print("Ciphertext is:")
    print(pt_grid)
    cyphertext = ""
    for block in range(len(pt_grid)):
        for col in range(4):
            for cell in range(4):
                temp = chr(pt_grid[block][col][cell])
                cyphertext = cyphertext + temp

    for block in range(len(pt_grid)):
        for col in range(4):
            for cell in range(4):
                pt_grid[block][col][cell] = hex(pt_grid[block][col][cell])
    print(pt_grid)
    print(cyphertext)
