local sha = require("sha")

local aes = {}

local 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
}

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

local RCON = {
    0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1B, 0x36
}

local function derive_key(secret)
    local _, bin = sha.sha256(secret)
    return bin:sub(17,32)
end

local function bxor(...)
    local args = {...}
    local result = args[1]
    for i = 2, #args do
        result = bit.bxor(result, args[i])
    end
    return result
end

local function xor(block1, block2)
    local result = {}
    for i = 1, #block1 do
        result[i] = string.char(bit.bxor(string.byte(block1, i), string.byte(block2, i)))
    end
    return table.concat(result)
end

local function to_matrix(plainText)
    local matrix = {}
    for i = 1, 4 do
        matrix[i] = {}
        for j = 1, 4 do
            matrix[i][j] = string.byte(plainText, (j - 1) * 4 + i)
        end
    end
    return matrix
end

local function from_matrix(matrix)
    local block = {}
    for i = 1, 4 do
        for j = 1, 4 do
            block[#block + 1] = string.char(matrix[j][i])
        end
    end

    return table.concat(block)
end


local function gmul(a, b)
    local return_value = 0
    local temp = 0
    while a ~= 0 do
        if bit.band(a, 1) ~= 0 then
            return_value = bit.bxor(return_value, b)
        end
        temp = bit.band(b, 0x80)
        b = bit.blshift(b, 1)
        if temp ~= 0 then
            b = bit.bxor(b, 0x1b)
        end
        a = bit.brshift(bit.band(a, 0xff), 1)
    end
    return bit.band(return_value, 0xff)
end

local function rot_word(word)
    return { word[2], word[3], word[4], word[1]}
end

local function sub_word(word)
    local newWord = {}
    for i = 1, 4 do
        local hex = string.byte(word[i]) + 1
        newWord[i] = S_BOX[hex]
    end
    return newWord
end

local function xor_word(wordA, wordB)
    local newWord = {}
    for i = 1, 4 do
        newWord[i] = bit.bxor(wordA[i], wordB[i])
    end
    return newWord
end

local function key_expansion(key)
    local round_keys = {}
    round_keys[1] = to_matrix(key)
    local rcon_index = 1

    for round = 2, 11 do
        round_keys[round] = {}
        local temp = round_keys[round - 1][4]
        temp = rot_word(temp)
        temp = sub_word(temp)
        temp[1] = bit.bxor(temp[1], RCON[rcon_index])
        rcon_index = rcon_index + 1
        local w_i_minus_4 = round_keys[round - 1][1]
        round_keys[round][1] = xor_word(w_i_minus_4, temp)

        for row = 2, 4 do
            temp = round_keys[round][row - 1]
            w_i_minus_4 = round_keys[round - 1][row]
            round_keys[round][row] = xor_word(w_i_minus_4, temp)
        end
    end

    return round_keys
end

local function sub_bytes(block, decrypt)
    local box = decrypt and INV_S_BOX or S_BOX

    local new_block = {}
    for i = 1, 4 do
        new_block[i] = {}
        for j = 1, 4 do
            new_block[i][j] = box[block[i][j] + 1]
        end
    end
    return new_block
end

local function shift_rows(block)
    block[2] = {block[2][2], block[2][3], block[2][4], block[2][1]}

    block[3] = {block[3][3], block[3][4], block[3][1], block[3][2]}

    block[4] = {block[4][4], block[4][1], block[4][2], block[4][3]}
end

local function inv_shift_rows(block)
    block[2] = {block[2][4], block[2][1], block[2][2], block[2][3]}

    block[3] = {block[3][3], block[3][4], block[3][1], block[3][2]}

    block[4] = {block[4][2], block[4][3], block[4][4], block[4][1]}
end

local function mix_columns(state, decrypt)
    local temp = {}
    for i = 1, 4 do
        local s0 = state[1][i]
        local s1 = state[2][i]
        local s2 = state[3][i]
        local s3 = state[4][i]

        if decrypt then
            temp[1] = bxor(gmul(0x0e, s0), gmul(0x0b, s1), gmul(0x0d, s2), gmul(0x09, s3))
            temp[2] = bxor(gmul(0x09, s0), gmul(0x0e, s1), gmul(0x0b, s2), gmul(0x0d, s3))
            temp[3] = bxor(gmul(0x0d, s0), gmul(0x09, s1), gmul(0x0e, s2), gmul(0x0b, s3))
            temp[4] = bxor(gmul(0x0b, s0), gmul(0x0d, s1), gmul(0x09, s2), gmul(0x0e, s3))
        else
            temp[1] = bxor(gmul(0x02, s0), gmul(0x03, s1), gmul(0x01, s2), gmul(0x01, s3))
            temp[2] = bxor(gmul(0x01, s0), gmul(0x02, s1), gmul(0x03, s2), gmul(0x01, s3))
            temp[3] = bxor(gmul(0x01, s0), gmul(0x01, s1), gmul(0x02, s2), gmul(0x03, s3))
            temp[4] = bxor(gmul(0x03, s0), gmul(0x01, s1), gmul(0x01, s2), gmul(0x02, s3))
        end

        for j = 1, 4 do
            state[j][i] = temp[j]
        end
    end
end

local function add_round_key(block, round_key)
    for i = 1, 4 do
        for j = 1, 4 do
            block[i][j] = bit.bxor(block[i][j], round_key[i][j])
        end
    end
end

local function encryptBlock(block, round_keys)
    block = to_matrix(block)
    add_round_key(block, round_keys[1])

    for round = 2, 11 do
        block = sub_bytes(block)
        shift_rows(block)

        if round < 11 then
            mix_columns(block)
        end

        add_round_key(block, round_keys[round])
    end
    return from_matrix(block)
end


local function decryptBlock(block, round_keys)
    block = to_matrix(block)
    add_round_key(block, round_keys[11])

    for round = 10, 1, -1 do
        inv_shift_rows(block)
        block = sub_bytes(block, true)

        add_round_key(block, round_keys[round])

        if round > 1 then
            mix_columns(block, true)
        end
    end
    return from_matrix(block)
end

function aes.cbc_encrypt(message, secret, iv)
    local key = derive_key(secret)
    key = key_expansion(key)

    if not iv or #iv ~= 16 then
        local temp_iv = {}
        for _ = 1, 16 do
            temp_iv[#temp_iv + 1] = string.char(math.random(0, 255))
        end
        iv = table.concat(temp_iv)
    end

    local remainder = #message % 16
    local fill_length = (remainder == 0) and 16 or (16 - remainder)
    message = message .. string.rep(string.char(fill_length), fill_length)

    local blocks = {}
    for i = 1, #message, 16 do
        table.insert(blocks, message:sub(i, i + 15))
    end

    local prev = iv
    local encrypted_message = ""

    for _, block in ipairs(blocks) do
        local xorBlock = xor(block, prev)
        local encrypted_block = encryptBlock(xorBlock, key)
        encrypted_message = encrypted_message .. encrypted_block
        prev = encrypted_block
    end

    return iv .. encrypted_message
end

function aes.cbc_decrypt(message, secret)
    local key = derive_key(secret)
    key = key_expansion(key)

    local iv = message:sub(1, 16)
    message = message:sub(17)

    local blocks = {}
    for i = 1, #message, 16 do
        table.insert(blocks, message:sub(i, i + 15))
    end

    local prev = iv
    local decrypted_message = ""

    for _, block in ipairs(blocks) do
        local decrypted_block = decryptBlock(block, key)
        local xorBlock = xor(decrypted_block, prev)
        decrypted_message = decrypted_message .. xorBlock
        prev = block
    end

    local fill_length = string.byte(decrypted_message:sub(-1))
    if fill_length >= 1 and fill_length <= 16 then
        decrypted_message = decrypted_message:sub(1, -fill_length - 1)
    end

    return decrypted_message
end

function aes.generate_iv()
    local t = {}
    for i = 1, 16 do t[i] = string.char(math.random(0,255)) end
    return table.concat(t)
end

return aes