/*

 * Copyright (c) 2020, Ali Mohammad Pur <mpfard@serenityos.org>

 * Copyright (c) 2022, the SerenityOS developers.

 *

 * SPDX-License-Identifier: BSD-2-Clause

 */

#pragma once

#include <AK/CPUFeatures.h>

#include <AK/Vector.h>

#include <LibCrypto/Cipher/Cipher.h>

#include <LibCrypto/Cipher/Mode/CBC.h>

#include <LibCrypto/Cipher/Mode/CTR.h>

#include <LibCrypto/Cipher/Mode/GCM.h>

#ifndef KERNEL

#    include <AK/ByteString.h>

#endif

namespace Crypto::Cipher {

struct AESCipherBlock : public CipherBlock {

public:

    static constexpr size_t BlockSizeInBits = 128;

    explicit AESCipherBlock(PaddingMode mode = PaddingMode::CMS)

        : CipherBlock(mode)

    {

    }

    AESCipherBlock(u8 const* data, size_t length, PaddingMode mode = PaddingMode::CMS)

        : AESCipherBlock(mode)

    {

        CipherBlock::overwrite(data, length);

    }

    constexpr static size_t block_size() { return BlockSizeInBits / 8; }

    virtual ReadonlyBytes bytes() const override { return ReadonlyBytes { m_data, sizeof(m_data) }; }

    virtual Bytes bytes() override { return Bytes { m_data, sizeof(m_data) }; }

    virtual void overwrite(ReadonlyBytes) override;

    virtual void overwrite(u8 const* data, size_t size) override { overwrite({ data, size }); }

    virtual void apply_initialization_vector(ReadonlyBytes ivec) override

    {

        for (size_t i = 0; i < min(block_size(), ivec.size()); ++i)

            m_data[i] ^= ivec[i];

    }

#ifndef KERNEL

    ByteString to_byte_string() const;

#endif

private:

    constexpr static size_t data_size() { return sizeof(m_data); }

    u8 m_data[BlockSizeInBits / 8] {};

};

struct AESCipherKey : public CipherKey {

    virtual ReadonlyBytes bytes() const override { return ReadonlyBytes { m_rd_keys, sizeof(m_rd_keys) }; }

    virtual void expand_encrypt_key(ReadonlyBytes user_key, size_t bits) override { return (this->*expand_encrypt_key_dispatched)(user_key, bits); }

    virtual void expand_decrypt_key(ReadonlyBytes user_key, size_t bits) override { return (this->*expand_decrypt_key_dispatched)(user_key, bits); }

    static bool is_valid_key_size(size_t bits) { return bits == 128 || bits == 192 || bits == 256; }

#ifndef KERNEL

    ByteString to_byte_string() const;

#endif

    u32 const* round_keys() const

    {

        return (u32 const*)m_rd_keys;

    }

    AESCipherKey(ReadonlyBytes user_key, size_t key_bits, Intent intent)

        : m_bits(key_bits)

    {

        if (intent == Intent::Encryption)

            expand_encrypt_key(user_key, key_bits);

        else

            expand_decrypt_key(user_key, key_bits);

    }

    virtual ~AESCipherKey() override = default;

    size_t rounds() const { return m_rounds; }

    size_t length() const { return m_bits / 8; }

protected:

    u32* round_keys()

    {

        return (u32*)m_rd_keys;

    }

private:

    template<CPUFeatures>

    void expand_encrypt_key_impl(ReadonlyBytes user_key, size_t bits);

    template<CPUFeatures>

    void expand_decrypt_key_impl(ReadonlyBytes user_key, size_t bits);

    static void (AESCipherKey::* const expand_encrypt_key_dispatched)(ReadonlyBytes user_key, size_t bits);

    static void (AESCipherKey::* const expand_decrypt_key_dispatched)(ReadonlyBytes user_key, size_t bits);

    static constexpr size_t MAX_ROUND_COUNT = 14;

    u32 m_rd_keys[(MAX_ROUND_COUNT + 1) * 4] { 0 };

    size_t m_rounds;

    size_t m_bits;

};

class AESCipher final : public Cipher<AESCipherKey, AESCipherBlock> {

public:

    using CBCMode = CBC<AESCipher>;

    using CTRMode = CTR<AESCipher>;

    using GCMMode = GCM<AESCipher>;

    constexpr static size_t BlockSizeInBits = BlockType::BlockSizeInBits;

    AESCipher(ReadonlyBytes user_key, size_t key_bits, Intent intent = Intent::Encryption, PaddingMode mode = PaddingMode::CMS)

        : Cipher<AESCipherKey, AESCipherBlock>(mode)

        , m_key(user_key, key_bits, intent)

    {

    }

    virtual AESCipherKey const& key() const override { return m_key; }

    virtual AESCipherKey& key() override { return m_key; }

    virtual void encrypt_block(BlockType const& in, BlockType& out) override { return (this->*encrypt_block_dispatched)(in, out); }

    virtual void decrypt_block(BlockType const& in, BlockType& out) override { return (this->*decrypt_block_dispatched)(in, out); }

#ifndef KERNEL

    virtual ByteString class_name() const override

    {

        return "AES";

    }

#endif

protected:

    AESCipherKey m_key;

private:

    template<CPUFeatures>

    void encrypt_block_impl(BlockType const& in, BlockType& out);

    template<CPUFeatures>

    void decrypt_block_impl(BlockType const& in, BlockType& out);

    static void (AESCipher::* const encrypt_block_dispatched)(BlockType const& in, BlockType& out);

    static void (AESCipher::* const decrypt_block_dispatched)(BlockType const& in, BlockType& out);

};

}