/**
 * @file aes.h
 * @brief AES-128/AES-256 block cipher with CBC, CTS, and AEAD-GCM modes.
 *
 * Declares netplus's self-contained AES implementation: a constant-time,
 * lookup-table-free software fallback plus an AES-NI/PCLMULQDQ
 * hardware-accelerated fast path for x86/x64 CPUs that support it (CPU
 * feature detection happens once at process startup in aes.cpp). Two
 * concrete key-size classes, aes128 and aes256, both implement the aes
 * interface and provide single-block ECB, CBC, ciphertext-stealing (CTS,
 * RFC 3962, used by Kerberos) and GCM authenticated encryption/decryption.
 * A small static helper class, aes_gcm, offers a cipher-agnostic entry
 * point into the GCM operations declared on aes.
 */
#pragma once
#include <vector>
#include <array>
#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <algorithm>

namespace netplus {

/** @brief A 16-byte AES block (plaintext or ciphertext), stored as a byte vector. */
using block128 = std::vector<uint8_t>;
/** @brief An AES key (16 bytes for AES-128, 32 bytes for AES-256), stored as a byte vector. */
using key128   = std::vector<uint8_t>;

/**
 * @brief Abstract interface for an AES block cipher instance bound to a fixed key.
 *
 * Declares the operations common to both key sizes (aes128, aes256): plain
 * block/CBC encryption and decryption, single-block ECB encryption (used for
 * QUIC/TLS header protection), and AES-GCM authenticated encryption. Concrete
 * subclasses hold the expanded round-key schedule and the GHASH subkey and
 * implement the actual AES rounds and GCM finite-field arithmetic, either via
 * AES-NI/PCLMULQDQ or a constant-time software fallback (see aes128 and
 * aes256).
 *
 * Instances are immutable after construction (the key schedule is fixed at
 * construction time), so concurrent encrypt/decrypt calls from multiple
 * threads on the same instance are safe as long as callers don't share
 * output buffers.
 */
class aes {
protected:
    // GCM operations - implemented by subclasses
    /**
     * @brief Multiplies the GHASH accumulator @p X by the subkey H in GF(2^128), in place.
     *
     * Implemented by each subclass either via the PCLMULQDQ carry-less
     * multiply instruction (hardware path) or a constant-time bit-by-bit
     * shift-and-reduce loop (software fallback). Used internally by ghash().
     * @param X 16-byte GHASH accumulator/operand, updated in place with the result.
     */
    virtual void gcm_mul(std::array<uint8_t,16>& X) const = 0;

    /**
     * @brief Computes the GHASH universal hash over the AAD and the ciphertext.
     *
     * Folds in the additional authenticated data, then the ciphertext, and
     * finally a 16-byte length block (bit-lengths of AAD and ciphertext,
     * big-endian) as specified by NIST SP 800-38D. The result is what
     * aes_gcm_encrypt()/aes_gcm_decrypt() XOR with E_K(J0) to produce or
     * verify the authentication tag.
     * @param out 16-byte output buffer receiving the GHASH result.
     * @param aad Pointer to the additional authenticated data (may be null if @p aad_len is 0).
     * @param aad_len Length of @p aad in bytes.
     * @param ct Pointer to the ciphertext bytes to authenticate.
     * @param ct_len Length of @p ct in bytes.
     */
    virtual void ghash(std::array<uint8_t,16>& out,
                       const uint8_t* aad, size_t aad_len,
                       const uint8_t* ct,  size_t ct_len) const = 0;

    // Allow aes_gcm to access protected methods
    /** @brief Grants aes_gcm access to the protected GCM primitives above. */
    friend class aes_gcm;

public:
    /** @brief Virtual destructor; concrete subclasses hold no external resources beyond their key schedule. */
    virtual ~aes() = default;

    /**
     * @brief Encrypts a single 16-byte block in raw (ECB) mode, no chaining.
     * @param plaintext Exactly one 16-byte AES block.
     * @return The 16-byte ciphertext block.
     * @throws std::invalid_argument if @p plaintext is not exactly 16 bytes.
     * @note ECB alone provides no semantic security (identical blocks
     *       produce identical output); this is a low-level primitive used
     *       internally by the CBC/CTS/GCM modes below and by callers that
     *       genuinely need single, independent block operations. Prefer
     *       encryptCBC() or aes_gcm_encrypt() for encrypting real data.
     */
    virtual std::vector<uint8_t> encrypt(const std::vector<uint8_t>& plaintext) = 0;

    /**
     * @brief Decrypts a single 16-byte block in raw (ECB) mode, no chaining.
     * @param ciphertext Exactly one 16-byte AES block.
     * @return The 16-byte plaintext block.
     * @throws std::invalid_argument if @p ciphertext is not exactly 16 bytes.
     */
    virtual std::vector<uint8_t> decrypt(const std::vector<uint8_t>& ciphertext) = 0;

    // Single-block ECB encrypt (no allocation) — for header protection
    /**
     * @brief Single-block ECB encryption into caller-supplied buffers (no heap allocation).
     *
     * Used for QUIC/TLS header protection, where a single AES block must be
     * encrypted on the hot path without the allocation overhead of the
     * std::vector-based encrypt().
     * @param in 16-byte input block.
     * @param out 16-byte output buffer.
     */
    virtual void encrypt_ecb(const uint8_t in[16], uint8_t out[16]) = 0;

    /**
     * @brief Encrypts data in CBC (Cipher Block Chaining) mode.
     * @param plaintext Plaintext whose length must be a non-zero multiple of 16 bytes.
     * @param iv 16-byte initialization vector.
     * @return Ciphertext of the same length as @p plaintext.
     * @throws std::invalid_argument if @p iv is not 16 bytes, or @p plaintext
     *         is empty or not a multiple of 16 bytes.
     */
    virtual std::vector<uint8_t> encryptCBC(const std::vector<uint8_t>& plaintext,
                                            const std::vector<uint8_t>& iv) = 0;
    /**
     * @brief Decrypts data in CBC (Cipher Block Chaining) mode.
     * @param ciphertext Ciphertext whose length must be a non-zero multiple of 16 bytes.
     * @param iv 16-byte initialization vector; must match the one used to encrypt.
     * @return Plaintext of the same length as @p ciphertext.
     * @throws std::invalid_argument if @p iv is not 16 bytes, or @p ciphertext
     *         is empty or not a multiple of 16 bytes.
     */
    virtual std::vector<uint8_t> decryptCBC(const std::vector<uint8_t>& ciphertext,
                                            const std::vector<uint8_t>& iv) = 0;

    // GCM encryption/decryption
    /**
     * @brief Encrypts and authenticates data using AES-GCM (NIST SP 800-38D).
     *
     * Computes ciphertext and a 16-byte authentication tag over @p pt and the
     * additional authenticated data @p aad, using J0 = @p iv12 || 0x00000001
     * as the initial GCM counter block. @p ct may alias @p pt for in-place
     * operation, since each output byte only depends on the corresponding
     * input byte and the CTR keystream.
     * @param iv12 12-byte nonce; must never be reused with the same key.
     * @param aad Additional authenticated data, not encrypted but covered by the tag.
     * @param aad_len Length of @p aad in bytes.
     * @param pt Plaintext to encrypt.
     * @param pt_len Length of @p pt (and of @p ct) in bytes.
     * @param ct Output buffer for the ciphertext, at least @p pt_len bytes.
     * @param tag16 Output buffer for the 16-byte authentication tag.
     * @return true on success (GCM encryption itself cannot otherwise fail).
     * @note Dispatches to a fused AES-NI + PCLMULQDQ CTR/GHASH fast path
     *       when the CPU supports it, falling back to a constant-time
     *       software implementation otherwise (see aes.cpp's runtime
     *       AES-NI/PCLMULQDQ CPUID detection).
     */
    virtual bool aes_gcm_encrypt(
        const uint8_t iv12[12],
        const uint8_t* aad, size_t aad_len,
        const uint8_t* pt,  size_t pt_len,
        uint8_t* ct,
        uint8_t tag16[16]
    ) const = 0;

    /**
     * @brief Decrypts and verifies AES-GCM ciphertext (NIST SP 800-38D).
     *
     * Recomputes the expected authentication tag and compares it against
     * @p tag16 in constant time before the plaintext is considered valid.
     * @param iv12 12-byte nonce; must match the value used at encryption.
     * @param aad Additional authenticated data covered by the tag.
     * @param aad_len Length of @p aad in bytes.
     * @param ct Ciphertext to decrypt.
     * @param ct_len Length of @p ct (and of @p pt) in bytes.
     * @param tag16 16-byte authentication tag to verify.
     * @param pt Output buffer for the plaintext, at least @p ct_len bytes.
     * @return true if authentication succeeded and @p pt holds the recovered
     *         plaintext; false if the tag does not match, in which case the
     *         AEAD contract is all-or-nothing and @p pt must not be used by
     *         the caller (implementations defensively zero it on failure).
     * @note Authentication failure means the ciphertext, AAD, key, nonce or
     *       tag do not all agree (e.g. tampering or corruption) — the data
     *       must be discarded, never processed.
     */
    virtual bool aes_gcm_decrypt(
        const uint8_t iv12[12],
        const uint8_t* aad, size_t aad_len,
        const uint8_t* ct,  size_t ct_len,
        const uint8_t tag16[16],
        uint8_t* pt
    ) const = 0;
};

/**
 * @brief AES-128 (10-round) block cipher with CBC, CTS, and GCM modes.
 *
 * Implements the aes interface for 128-bit keys. All round transformations
 * (SubBytes/InvSubBytes, ShiftRows, MixColumns, key schedule S-box) use a
 * constant-time GF(2^8) implementation with no lookup tables, avoiding
 * cache-timing side channels in the software fallback path. When the host
 * CPU exposes AES-NI and PCLMULQDQ (detected once at process startup via
 * CPUID), encrypt_block()/decrypt_block() and the GCM CTR/GHASH loops
 * dispatch to hand-written intrinics that use the hardware AES round
 * instructions and carry-less multiplication instead.
 *
 * @see aes256 for the 256-bit-key counterpart; the two classes are
 *      near-identical except for round count (10 vs. 14) and key-schedule
 *      size/shape.
 */
class aes128 final : public aes {
private:
    /** @brief Expanded AES-128 round-key schedule: 11 round keys x 16 bytes = 176 bytes, populated by keyExpansion(). */
    std::vector<uint8_t> m_roundKeyBytes; // 176 bytes for AES-128
    /** @brief GHASH subkey H = E_K(0^128), precomputed once at construction time and used by gcm_mul()/ghash(). */
    uint8_t m_H[16];                      // GHASH subkey

    /** @brief Returns a pointer to the 16-byte round key for round @p round within m_roundKeyBytes. */
    inline const uint8_t* rk(int round) const {
        return &m_roundKeyBytes[size_t(round) * 16];
    }

    // ---- Constant-time helper ops ----
    /** @brief Constant-time byte equality test: returns 0x01 if @p a == @p b, else 0x00. */
    static inline uint8_t ct_eq_u8(uint8_t a, uint8_t b) {
        uint8_t x = static_cast<uint8_t>(a ^ b);
        x |= x >> 4; x |= x >> 2; x |= x >> 1;
        return static_cast<uint8_t>((x ^ 1) & 1);
    }

    /** @brief Constant-time byte select: returns @p a if @p mask is 0xFF, @p b if @p mask is 0x00. */
    static inline uint8_t ct_select_u8(uint8_t mask, uint8_t a, uint8_t b) {
        // mask: 0x00 or 0xFF
        return static_cast<uint8_t>((a & mask) | (b & ~mask));
    }

    /** @brief Rotates an 8-bit value left by @p r bits. */
    static inline uint8_t rotl8(uint8_t x, unsigned r) {
        return static_cast<uint8_t>((x << r) | (x >> (8 - r)));
    }

    // ---- GF(2^8) multiply for AES S-box ----
    /** @brief Constant-time multiplication of @p a and @p b in GF(2^8) under the AES reduction polynomial (0x11B). */
    static uint8_t gf_mul(uint8_t a, uint8_t b);
    /** @brief Constant-time multiplicative inverse of @p a in GF(2^8) (computed via exponentiation to the 254th power; 0 maps to 0). */
    static uint8_t gf_inv(uint8_t a);

    // ---- constant-time AES S-box + inv S-box ----
    /** @brief Constant-time AES S-box: GF(2^8) inverse followed by the AES affine transform (no lookup table). */
    static uint8_t ct_sbox(uint8_t x);
    /** @brief Constant-time AES inverse S-box: inverse affine transform followed by GF(2^8) inverse. */
    static uint8_t ct_inv_sbox(uint8_t x);

    // ---- key schedule ----
    /** @brief Rotates a 32-bit key-schedule word left by 8 bits (the AES RotWord step). */
    static uint32_t rotWord(uint32_t w);
    /** @brief Applies the constant-time S-box to each byte of a 32-bit key-schedule word (the AES SubWord step). */
    static uint32_t subWord_ct(uint32_t w);
    /**
     * @brief Expands a 16-byte AES-128 key into the 176-byte, 11-round key schedule stored in m_roundKeyBytes.
     * @param key The AES-128 key; must be exactly 16 bytes.
     * @throws std::invalid_argument if @p key is not 16 bytes.
     */
    void keyExpansion(const key128& key);

    // ---- AES core block ops ----
    /**
     * @brief Encrypts one 16-byte block through all 10 AES-128 rounds.
     * @param in 16-byte plaintext block.
     * @param out 16-byte ciphertext output buffer.
     * @note Uses the AES-NI hardware path when available (g_has_hw_aes),
     *       otherwise the constant-time software round functions below.
     */
    void encrypt_block(const uint8_t in[16], uint8_t out[16]) const;
    /**
     * @brief Decrypts one 16-byte block through all 10 AES-128 inverse rounds.
     * @param in 16-byte ciphertext block.
     * @param out 16-byte plaintext output buffer.
     * @note Uses the AES-NI hardware path when available (g_has_hw_aes),
     *       otherwise the constant-time software round functions below.
     */
    void decrypt_block(const uint8_t in[16], uint8_t out[16]) const;

    /** @brief XORs the 16-byte round key @p roundKey into the state @p state (AddRoundKey). */
    static void addRoundKey(uint8_t* state, const uint8_t* roundKey);
    /** @brief Applies the constant-time S-box to every byte of the state (SubBytes). */
    static void subBytes(uint8_t* state);
    /** @brief Applies the constant-time inverse S-box to every byte of the state (InvSubBytes). */
    static void invSubBytes(uint8_t* state);
    /** @brief Cyclically shifts each row of the column-major state left by its row index (ShiftRows). */
    static void shiftRows(uint8_t* state);
    /** @brief Inverse of shiftRows(): cyclically shifts each row right by its row index (InvShiftRows). */
    static void invShiftRows(uint8_t* state);
    /** @brief Mixes each column of the state via GF(2^8) matrix multiplication (MixColumns). */
    static void mixColumns(uint8_t* state);
    /** @brief Inverse of mixColumns(): un-mixes each column via the inverse GF(2^8) matrix (InvMixColumns). */
    static void invMixColumns(uint8_t* state);

    /** @brief GF(2^8) "multiply by x" (xtime): shifts left by one bit, reducing modulo the AES polynomial (0x1B) if the top bit was set. */
    static inline uint8_t xtime(uint8_t x) {
        return static_cast<uint8_t>((x << 1) ^ ((x & 0x80) ? 0x1B : 0x00));
    }

    /** @brief GF(2^8) multiplication of @p a and @p b, used by invMixColumns() for its larger fixed coefficients. */
    static uint8_t gmul(uint8_t a, uint8_t b);

    // ---- GCM helpers ----
    /** @brief Increments the 32-bit big-endian counter held in the last 4 bytes of a 16-byte GCM counter block, in place. */
    static inline void inc32(uint8_t ctr[16]) {
        uint32_t x = (uint32_t(ctr[12]) << 24) |
                    (uint32_t(ctr[13]) << 16) |
                    (uint32_t(ctr[14]) << 8)  |
                    (uint32_t(ctr[15]));
        x++;
        ctr[12] = uint8_t(x >> 24);
        ctr[13] = uint8_t(x >> 16);
        ctr[14] = uint8_t(x >> 8);
        ctr[15] = uint8_t(x);
    }

    /** @brief Loads 8 bytes at @p p as a big-endian uint64_t. */
    static inline uint64_t load_be64(const uint8_t* p) {
        return (uint64_t(p[0]) << 56) |
               (uint64_t(p[1]) << 48) |
               (uint64_t(p[2]) << 40) |
               (uint64_t(p[3]) << 32) |
               (uint64_t(p[4]) << 24) |
               (uint64_t(p[5]) << 16) |
               (uint64_t(p[6]) << 8)  |
               (uint64_t(p[7]));
    }

    /** @brief Stores @p x into 8 bytes at @p p in big-endian order. */
    static inline void store_be64(uint8_t* p, uint64_t x) {
        p[0] = uint8_t(x >> 56);
        p[1] = uint8_t(x >> 48);
        p[2] = uint8_t(x >> 40);
        p[3] = uint8_t(x >> 32);
        p[4] = uint8_t(x >> 24);
        p[5] = uint8_t(x >> 16);
        p[6] = uint8_t(x >> 8);
        p[7] = uint8_t(x);
    }

    /**
     * @brief Constant-time comparison of two @p n-byte buffers.
     *
     * Used to compare the computed GCM authentication tag against the
     * wire tag without leaking timing information about where the first
     * mismatching byte occurs.
     * @param a First buffer.
     * @param b Second buffer.
     * @param n Number of bytes to compare.
     * @return true if all @p n bytes are equal.
     */
    static inline bool ct_memeq(const uint8_t* a, const uint8_t* b, size_t n) {
        uint8_t diff = 0;
        for (size_t i = 0; i < n; i++)
            diff |= static_cast<uint8_t>(a[i] ^ b[i]);
        return diff == 0;
    }

protected:
    /** @brief AES-128 GHASH field multiply; see aes::gcm_mul(). Uses PCLMULQDQ when available, otherwise a constant-time bit loop. */
    void gcm_mul(std::array<uint8_t,16>& X) const override;
    /** @brief AES-128 GHASH computation; see aes::ghash(). */
    void ghash(std::array<uint8_t,16>& out,
               const uint8_t* aad, size_t aad_len,
               const uint8_t* ct,  size_t ct_len) const override;

public:
    /**
     * @brief Constructs an AES-128 cipher instance and expands the key schedule.
     * @param key The AES-128 key; must be exactly 16 bytes.
     * @throws std::invalid_argument if @p key is not 16 bytes (via keyExpansion()).
     */
    explicit aes128(const key128& key);
    /** @brief Destructor (no explicit cleanup; key material is held in a std::vector). */
    ~aes128() = default;

    /** @copydoc aes::encrypt */
    std::vector<uint8_t> encrypt(const std::vector<uint8_t>& plaintext) override;
    /** @copydoc aes::decrypt */
    std::vector<uint8_t> decrypt(const std::vector<uint8_t>& ciphertext) override;
    /** @copydoc aes::encrypt_ecb */
    void encrypt_ecb(const uint8_t in[16], uint8_t out[16]) override { encrypt_block(in, out); }

    /** @copydoc aes::encryptCBC */
    std::vector<uint8_t> encryptCBC(const std::vector<uint8_t>& plaintext,
                                    const std::vector<uint8_t>& iv) override;

    /** @copydoc aes::decryptCBC */
    std::vector<uint8_t> decryptCBC(const std::vector<uint8_t>& ciphertext,
                                    const std::vector<uint8_t>& iv) override;

    // CTS mode (RFC 3962 - Kerberos AES encryption)
    // Input may be any length >= 16 bytes (one AES block).
    // No padding is needed; output size equals input size.
    /**
     * @brief Encrypts data using CBC-CS3 ciphertext stealing (RFC 3962, as used by Kerberos).
     *
     * Unlike plain CBC, the input may be any length that is at least one
     * full block (16 bytes) — no padding to a block multiple is required,
     * and the output is exactly the same size as the input. Internally,
     * the plaintext is zero-padded to a block boundary, CBC-encrypted, and
     * then the last two ciphertext blocks are swapped and the final block
     * truncated to the true plaintext length (CBC-CS3 variant).
     * @param plaintext Plaintext of at least 16 bytes; any length >= 16 is accepted.
     * @param iv 16-byte initialization vector.
     * @return Ciphertext of exactly the same length as @p plaintext.
     * @throws std::invalid_argument if @p iv is not 16 bytes or @p plaintext is shorter than 16 bytes.
     * @see decryptCTS
     */
    std::vector<uint8_t> encryptCTS(const std::vector<uint8_t>& plaintext,
                                    const std::vector<uint8_t>& iv);

    /**
     * @brief Decrypts CBC-CS3 ciphertext-stolen data (RFC 3962, as used by Kerberos).
     * @param ciphertext Ciphertext of at least 16 bytes, as produced by encryptCTS().
     * @param iv 16-byte initialization vector; must match the one used to encrypt.
     * @return Plaintext of exactly the same length as @p ciphertext.
     * @throws std::invalid_argument if @p iv is not 16 bytes or @p ciphertext is shorter than 16 bytes.
     * @see encryptCTS
     */
    std::vector<uint8_t> decryptCTS(const std::vector<uint8_t>& ciphertext,
                                    const std::vector<uint8_t>& iv);

    /** @copydoc aes::aes_gcm_encrypt */
    bool aes_gcm_encrypt(
        const uint8_t iv12[12],
        const uint8_t* aad, size_t aad_len,
        const uint8_t* pt,  size_t pt_len,
        uint8_t* ct,
        uint8_t tag16[16]
    ) const  override;

    /** @copydoc aes::aes_gcm_decrypt */
    bool aes_gcm_decrypt(
        const uint8_t iv12[12],
        const uint8_t* aad, size_t aad_len,
        const uint8_t* ct,  size_t ct_len,
        const uint8_t tag16[16],
        uint8_t* pt
    ) const override;
};

/**
 * @brief AES-256 (14-round) block cipher with CBC, CTS, and GCM modes.
 *
 * Implements the aes interface for 256-bit keys. Structurally identical to
 * aes128 (same constant-time software S-box/round functions and the same
 * AES-NI/PCLMULQDQ hardware dispatch), differing only in round count (14
 * vs. 10), round-key schedule size (240 vs. 176 bytes), and the AES-256 key
 * schedule's alternating SubWord application (see keyExpansion()).
 *
 * @see aes128 for full documentation of the shared design (constant-time
 *      S-box, AES-NI dispatch, GCM tag handling); only AES-256-specific
 *      details are repeated here.
 */
class aes256 final : public aes {
private:
    /** @brief Expanded AES-256 round-key schedule: 15 round keys x 16 bytes = 240 bytes, populated by keyExpansion(). */
    std::vector<uint8_t> m_roundKeyBytes; // 240 bytes for AES-256
    /** @brief GHASH subkey H = E_K(0^128), precomputed once at construction time and used by gcm_mul()/ghash(). */
    uint8_t m_H[16];                      // GHASH subkey

    /** @brief Returns a pointer to the 16-byte round key for round @p round within m_roundKeyBytes. */
    inline const uint8_t* rk(int round) const {
        return &m_roundKeyBytes[size_t(round) * 16];
    }

    /** @copydoc aes128::gf_mul */
    static uint8_t gf_mul(uint8_t a, uint8_t b);
    /** @copydoc aes128::gf_inv */
    static uint8_t gf_inv(uint8_t a);
    /** @copydoc aes128::ct_sbox */
    static uint8_t ct_sbox(uint8_t x);
    /** @copydoc aes128::ct_inv_sbox */
    static uint8_t ct_inv_sbox(uint8_t x);

    /** @copydoc aes128::rotWord */
    static uint32_t rotWord(uint32_t w);
    /** @copydoc aes128::subWord_ct */
    static uint32_t subWord_ct(uint32_t w);
    /**
     * @brief Expands a 32-byte AES-256 key into the 240-byte, 15-round key schedule stored in m_roundKeyBytes.
     *
     * Unlike AES-128, the AES-256 key schedule applies SubWord not only
     * every 8 words (RCON-XORed, as in AES-128) but also once more halfway
     * between (every 4th word when it is not already an RCON word) — see
     * the "Deliberately always the constant-time software S-box" note in
     * aes.cpp explaining why this always uses subWord_ct() rather than an
     * AES-NI keygen intrinsic even on hardware-accelerated builds.
     * @param key The AES-256 key; must be exactly 32 bytes.
     * @throws std::invalid_argument if @p key is not 32 bytes.
     */
    void keyExpansion(const std::vector<uint8_t>& key);

    /** @brief Encrypts one 16-byte block through all 14 AES-256 rounds. @see aes128::encrypt_block */
    void encrypt_block(const uint8_t in[16], uint8_t out[16]) const;
    /** @brief Decrypts one 16-byte block through all 14 AES-256 inverse rounds. @see aes128::decrypt_block */
    void decrypt_block(const uint8_t in[16], uint8_t out[16]) const;

    /** @copydoc aes128::addRoundKey */
    static void addRoundKey(uint8_t* state, const uint8_t* roundKey);
    /** @copydoc aes128::subBytes */
    static void subBytes(uint8_t* state);
    /** @copydoc aes128::invSubBytes */
    static void invSubBytes(uint8_t* state);
    /** @copydoc aes128::shiftRows */
    static void shiftRows(uint8_t* state);
    /** @copydoc aes128::invShiftRows */
    static void invShiftRows(uint8_t* state);
    /** @copydoc aes128::mixColumns */
    static void mixColumns(uint8_t* state);
    /** @copydoc aes128::invMixColumns */
    static void invMixColumns(uint8_t* state);

    /** @copydoc aes128::xtime */
    static inline uint8_t xtime(uint8_t x) {
        return static_cast<uint8_t>((x << 1) ^ ((x & 0x80) ? 0x1B : 0x00));
    }

    /** @copydoc aes128::gmul */
    static uint8_t gmul(uint8_t a, uint8_t b);

    /** @copydoc aes128::inc32 */
    static inline void inc32(uint8_t ctr[16]) {
        uint32_t x = (uint32_t(ctr[12]) << 24) |
                    (uint32_t(ctr[13]) << 16) |
                    (uint32_t(ctr[14]) << 8)  |
                    (uint32_t(ctr[15]));
        x++;
        ctr[12] = uint8_t(x >> 24);
        ctr[13] = uint8_t(x >> 16);
        ctr[14] = uint8_t(x >> 8);
        ctr[15] = uint8_t(x);
    }

    /** @copydoc aes128::load_be64 */
    static inline uint64_t load_be64(const uint8_t* p) {
        return (uint64_t(p[0]) << 56) |
               (uint64_t(p[1]) << 48) |
               (uint64_t(p[2]) << 40) |
               (uint64_t(p[3]) << 32) |
               (uint64_t(p[4]) << 24) |
               (uint64_t(p[5]) << 16) |
               (uint64_t(p[6]) << 8)  |
               (uint64_t(p[7]));
    }

    /** @copydoc aes128::store_be64 */
    static inline void store_be64(uint8_t* p, uint64_t x) {
        p[0] = uint8_t(x >> 56);
        p[1] = uint8_t(x >> 48);
        p[2] = uint8_t(x >> 40);
        p[3] = uint8_t(x >> 32);
        p[4] = uint8_t(x >> 24);
        p[5] = uint8_t(x >> 16);
        p[6] = uint8_t(x >> 8);
        p[7] = uint8_t(x);
    }

    /** @copydoc aes128::ct_memeq */
    static inline bool ct_memeq(const uint8_t* a, const uint8_t* b, size_t n) {
        uint8_t diff = 0;
        for (size_t i = 0; i < n; i++)
            diff |= static_cast<uint8_t>(a[i] ^ b[i]);
        return diff == 0;
    }

protected:
    /** @brief AES-256 GHASH field multiply; see aes::gcm_mul(). Uses PCLMULQDQ when available, otherwise a constant-time bit loop. */
    void gcm_mul(std::array<uint8_t,16>& X) const override;
    /** @brief AES-256 GHASH computation; see aes::ghash(). */
    void ghash(std::array<uint8_t,16>& out,
               const uint8_t* aad, size_t aad_len,
               const uint8_t* ct,  size_t ct_len) const override;

public:
    /**
     * @brief Constructs an AES-256 cipher instance and expands the key schedule.
     * @param key The AES-256 key; must be exactly 32 bytes.
     * @throws std::invalid_argument if @p key is not 32 bytes (via keyExpansion()).
     */
    explicit aes256(const std::vector<uint8_t>& key);
    /** @brief Destructor (no explicit cleanup; key material is held in a std::vector). */
    ~aes256() = default;

    /** @copydoc aes::encrypt */
    std::vector<uint8_t> encrypt(const std::vector<uint8_t>& plaintext) override;
    /** @copydoc aes::decrypt */
    std::vector<uint8_t> decrypt(const std::vector<uint8_t>& ciphertext) override;
    /** @copydoc aes::encrypt_ecb */
    void encrypt_ecb(const uint8_t in[16], uint8_t out[16]) override { encrypt_block(in, out); }

    /** @copydoc aes::encryptCBC */
    std::vector<uint8_t> encryptCBC(const std::vector<uint8_t>& plaintext,
                                    const std::vector<uint8_t>& iv) override;

    /** @copydoc aes::decryptCBC */
    std::vector<uint8_t> decryptCBC(const std::vector<uint8_t>& ciphertext,
                                    const std::vector<uint8_t>& iv) override;

    // CTS mode (RFC 3962 - Kerberos AES encryption)
    /** @copydoc aes128::encryptCTS */
    std::vector<uint8_t> encryptCTS(const std::vector<uint8_t>& plaintext,
                                    const std::vector<uint8_t>& iv);

    /** @copydoc aes128::decryptCTS */
    std::vector<uint8_t> decryptCTS(const std::vector<uint8_t>& ciphertext,
                                    const std::vector<uint8_t>& iv);

    /** @copydoc aes::aes_gcm_encrypt */
    bool aes_gcm_encrypt(
        const uint8_t iv12[12],
        const uint8_t* aad, size_t aad_len,
        const uint8_t* pt,  size_t pt_len,
        uint8_t* ct,
        uint8_t tag16[16]
    ) const override;

    /** @copydoc aes::aes_gcm_decrypt */
    bool aes_gcm_decrypt(
        const uint8_t iv12[12],
        const uint8_t* aad, size_t aad_len,
        const uint8_t* ct,  size_t ct_len,
        const uint8_t tag16[16],
        uint8_t* pt
    ) const override;
};

/**
 * @brief Cipher-agnostic static entry point for AES-GCM authenticated encryption.
 *
 * Thin static wrapper that forwards to the protected/virtual
 * aes::aes_gcm_encrypt()/aes::aes_gcm_decrypt() on a caller-supplied aes
 * instance (aes128 or aes256), so calling code can perform GCM operations
 * through a single, non-member API without needing to know the concrete
 * cipher type. Declared a friend of aes so it can reach those methods.
 */
class aes_gcm {
public:
    /**
     * @brief Encrypts and authenticates data with AES-GCM using an existing cipher instance.
     * @param cipher An initialized aes128 or aes256 instance providing the key.
     * @param iv12 12-byte nonce; must never be reused with the same key.
     * @param aad Additional authenticated data, not encrypted but covered by the tag.
     * @param aad_len Length of @p aad in bytes.
     * @param pt Plaintext to encrypt.
     * @param pt_len Length of @p pt (and of @p ct) in bytes.
     * @param ct Output buffer for the ciphertext, at least @p pt_len bytes.
     * @param tag16 Output buffer for the 16-byte authentication tag.
     * @return true on success.
     * @see aes::aes_gcm_encrypt
     */
    static bool encrypt(
        const aes& cipher,
        const uint8_t iv12[12],
        const uint8_t* aad, size_t aad_len,
        const uint8_t* pt,  size_t pt_len,
        uint8_t* ct,
        uint8_t tag16[16]
    );

    /**
     * @brief Decrypts and verifies AES-GCM ciphertext using an existing cipher instance.
     * @param cipher An initialized aes128 or aes256 instance providing the key.
     * @param iv12 12-byte nonce; must match the value used at encryption.
     * @param aad Additional authenticated data covered by the tag.
     * @param aad_len Length of @p aad in bytes.
     * @param ct Ciphertext to decrypt.
     * @param ct_len Length of @p ct (and of @p pt) in bytes.
     * @param tag16 16-byte authentication tag to verify.
     * @param pt Output buffer for the plaintext, at least @p ct_len bytes.
     * @return true if authentication succeeded and @p pt holds valid
     *         plaintext; false if the tag did not match, in which case
     *         @p pt must not be used.
     * @see aes::aes_gcm_decrypt
     */
    static bool decrypt(
        const aes& cipher,
        const uint8_t iv12[12],
        const uint8_t* aad, size_t aad_len,
        const uint8_t* ct,  size_t ct_len,
        const uint8_t tag16[16],
        uint8_t* pt
    );
};


} // namespace netplus

