#
#  Cipher/AES.py : AES
#
# ===================================================================
# The contents of this file are dedicated to the public domain.  To
# the extent that dedication to the public domain is not available,
# everyone is granted a worldwide, perpetual, royalty-free,
# non-exclusive license to exercise all rights associated with the
# contents of this file for any purpose whatsoever.
# No rights are reserved.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# ===================================================================

from __future__ import annotations

import sys
from typing import TYPE_CHECKING, Any, Literal, Optional, Union, overload

from Crypto.Cipher import _create_cipher
from Crypto.Random import get_random_bytes
from Crypto.Util import _cpu_features
from Crypto.Util._raw_api import (
    SmartPointer,
    VoidPointer,
    c_size_t,
    c_uint8_ptr_len,
    load_pycryptodome_raw_lib,
)

if TYPE_CHECKING:
    from Crypto.Cipher._mode_cbc import CbcMode
    from Crypto.Cipher._mode_ccm import CcmMode
    from Crypto.Cipher._mode_cfb import CfbMode
    from Crypto.Cipher._mode_ctr import CtrMode
    from Crypto.Cipher._mode_eax import EaxMode
    from Crypto.Cipher._mode_ecb import EcbMode
    from Crypto.Cipher._mode_gcm import GcmMode
    from Crypto.Cipher._mode_kw import KWMode
    from Crypto.Cipher._mode_kwp import KWPMode
    from Crypto.Cipher._mode_ocb import OcbMode
    from Crypto.Cipher._mode_ofb import OfbMode
    from Crypto.Cipher._mode_openpgp import OpenPgpMode
    from Crypto.Cipher._mode_siv import SivMode
    from Crypto.Util.Counter import CounterParams

Buffer = Union[bytes, bytearray, memoryview]

MODE_ECB: Literal[1] = 1  #: Electronic Code Book (:ref:`ecb_mode`)
MODE_CBC: Literal[2] = 2  #: Cipher-Block Chaining (:ref:`cbc_mode`)
MODE_CFB: Literal[3] = 3  #: Cipher Feedback (:ref:`cfb_mode`)
MODE_OFB: Literal[5] = 5  #: Output Feedback (:ref:`ofb_mode`)
MODE_CTR: Literal[6] = 6  #: Counter mode (:ref:`ctr_mode`)
MODE_OPENPGP: Literal[7] = 7  #: OpenPGP mode (:ref:`openpgp_mode`)
MODE_CCM: Literal[8] = 8  #: Counter with CBC-MAC (:ref:`ccm_mode`)
MODE_EAX: Literal[9] = 9  #: :ref:`eax_mode`
MODE_SIV: Literal[10] = 10  #: Synthetic Initialization Vector (:ref:`siv_mode`)
MODE_GCM: Literal[11] = 11  #: Galois Counter Mode (:ref:`gcm_mode`)
MODE_OCB: Literal[12] = 12  #: Offset Code Book (:ref:`ocb_mode`)
MODE_KW: Literal[13] = 13  #: Key Wrap (:ref:`kw_mode`)
MODE_KWP: Literal[14] = 14  #: Key Wrap with Padding (:ref:`kwp_mode`)

_cproto = """
        int AES_start_operation(const uint8_t key[],
                                size_t key_len,
                                void **pResult);
        int AES_encrypt(const void *state,
                        const uint8_t *in,
                        uint8_t *out,
                        size_t data_len);
        int AES_decrypt(const void *state,
                        const uint8_t *in,
                        uint8_t *out,
                        size_t data_len);
        int AES_stop_operation(void *state);
        """


# Load portable AES
_raw_aes_lib = load_pycryptodome_raw_lib("Crypto.Cipher._raw_aes", _cproto)


def _load_aes_hw():
    """Return the (start, stop) functions of the AES implementation
    with the AES instructions of the CPU (AES-NI on x86,
    the ARMv8 Cryptography Extensions on AArch64),
    or None if the CPU does not have them."""

    if _cpu_features.have_aes_ni():
        name, prefix = "Crypto.Cipher._raw_aesni", "AESNI"
    elif _cpu_features.have_aes_armv8():
        name, prefix = "Crypto.Cipher._raw_aes_armv8", "AES_ARMV8"
    else:
        return None
    try:
        lib = load_pycryptodome_raw_lib(name, _cproto.replace("AES", prefix))
    # The module may not have been compiled in
    except OSError:
        return None
    return (getattr(lib, prefix + "_start_operation"), getattr(lib, prefix + "_stop_operation"))


_raw_aes_hw = _load_aes_hw()


def _create_base_cipher(dict_parameters):
    """This method instantiates and returns a handle to a low-level
    base cipher. It will absorb named parameters in the process."""

    use_aesni = dict_parameters.pop("use_aesni", True)

    try:
        key = dict_parameters.pop("key")
    except KeyError:
        raise TypeError("Missing 'key' parameter")

    if len(key) not in key_size:
        raise ValueError("Incorrect AES key length (%d bytes)" % len(key))

    if use_aesni and _raw_aes_hw:
        start_operation, stop_operation = _raw_aes_hw
    else:
        start_operation = _raw_aes_lib.AES_start_operation
        stop_operation = _raw_aes_lib.AES_stop_operation

    cipher = VoidPointer()
    key_ptr, key_len = c_uint8_ptr_len(key)
    result = start_operation(key_ptr, c_size_t(key_len), cipher.address_of())
    if result:
        raise ValueError("Error %X while instantiating the AES cipher" % result)
    return SmartPointer(cipher.get(), stop_operation)


def _derive_Poly1305_key_pair(key, nonce):
    """Derive a tuple (r, s, nonce) for a Poly1305 MAC.

    If nonce is ``None``, a new 16-byte nonce is generated.
    """

    if len(key) != 32:
        raise ValueError("Poly1305 with AES requires a 32-byte key")

    if nonce is None:
        nonce = get_random_bytes(16)
    elif len(nonce) != 16:
        raise ValueError("Poly1305 with AES requires a 16-byte nonce")

    s = new(key[:16], MODE_ECB).encrypt(nonce)
    return key[16:], s, nonce


@overload
def new(key: Buffer, mode: Literal[1], use_aesni: bool = ...) -> EcbMode: ...


@overload
def new(key: Buffer, mode: Literal[2], iv: Optional[Buffer] = ..., use_aesni: bool = ...) -> CbcMode: ...


@overload
def new(key: Buffer, mode: Literal[2], IV: Optional[Buffer] = ..., use_aesni: bool = ...) -> CbcMode: ...


@overload
def new(
    key: Buffer, mode: Literal[3], iv: Optional[Buffer] = ..., segment_size: int = ..., use_aesni: bool = ...
) -> CfbMode: ...


@overload
def new(
    key: Buffer, mode: Literal[3], IV: Optional[Buffer] = ..., segment_size: int = ..., use_aesni: bool = ...
) -> CfbMode: ...


@overload
def new(key: Buffer, mode: Literal[5], iv: Optional[Buffer] = ..., use_aesni: bool = ...) -> OfbMode: ...


@overload
def new(key: Buffer, mode: Literal[5], IV: Optional[Buffer] = ..., use_aesni: bool = ...) -> OfbMode: ...


@overload
def new(
    key: Buffer,
    mode: Literal[6],
    nonce: Optional[Buffer] = ...,
    initial_value: Union[int, Buffer] = ...,
    counter: CounterParams = ...,
    use_aesni: bool = ...,
    threads: int = ...,
) -> CtrMode: ...


@overload
def new(key: Buffer, mode: Literal[7], iv: Optional[Buffer] = ..., use_aesni: bool = ...) -> OpenPgpMode: ...


@overload
def new(key: Buffer, mode: Literal[7], IV: Optional[Buffer] = ..., use_aesni: bool = ...) -> OpenPgpMode: ...


@overload
def new(
    key: Buffer,
    mode: Literal[8],
    nonce: Optional[Buffer] = ...,
    mac_len: int = ...,
    assoc_len: int = ...,
    use_aesni: bool = ...,
) -> CcmMode: ...


@overload
def new(
    key: Buffer, mode: Literal[9], nonce: Optional[Buffer] = ..., mac_len: int = ..., use_aesni: bool = ...
) -> EaxMode: ...


@overload
def new(key: Buffer, mode: Literal[10], nonce: Optional[Buffer] = ..., use_aesni: bool = ...) -> SivMode: ...


@overload
def new(
    key: Buffer,
    mode: Literal[11],
    nonce: Optional[Buffer] = ...,
    mac_len: int = ...,
    use_aesni: bool = ...,
    threads: int = ...,
) -> GcmMode: ...


@overload
def new(
    key: Buffer,
    mode: Literal[12],
    nonce: Optional[Buffer] = ...,
    mac_len: int = ...,
    use_aesni: bool = ...,
    threads: int = ...,
) -> OcbMode: ...


@overload
def new(key: Buffer, mode: Literal[13], use_aesni: bool = ...) -> KWMode: ...


@overload
def new(key: Buffer, mode: Literal[14], use_aesni: bool = ...) -> KWPMode: ...


def new(key: Buffer, mode: int, *args: Any, **kwargs: Any) -> Any:
    """Create a new AES cipher.

    Args:
      key(bytes/bytearray/memoryview):
        The secret key to use in the symmetric cipher.

        It must be 16 (*AES-128)*, 24 (*AES-192*) or 32 (*AES-256*) bytes long.

        For ``MODE_SIV`` only, it doubles to 32, 48, or 64 bytes.
      mode (a ``MODE_*`` constant):
        The chaining mode to use for encryption or decryption.
        If in doubt, use ``MODE_GCM`` (with a random 12-byte ``nonce``).

    Keyword Args:
      iv (bytes/bytearray/memoryview):
        (Only applicable for ``MODE_CBC``, ``MODE_CFB``, ``MODE_OFB``,
        and ``MODE_OPENPGP`` modes).

        The initialization vector to use for encryption or decryption.

        For ``MODE_CBC``, ``MODE_CFB``, and ``MODE_OFB`` it must be 16 bytes long.

        For ``MODE_OPENPGP`` mode only,
        it must be 16 bytes long for encryption
        and 18 bytes for decryption (in the latter case, it is
        actually the *encrypted* IV which was prefixed to the ciphertext).

        If not provided, a random byte string is generated (you must then
        read its value with the :attr:`iv` attribute).

      nonce (bytes/bytearray/memoryview):
        (Only applicable for ``MODE_CCM``, ``MODE_EAX``, ``MODE_GCM``,
        ``MODE_SIV``, ``MODE_OCB``, and ``MODE_CTR``).

        A value that must never be reused for any other encryption done
        with this key (except possibly for ``MODE_SIV``, see below).

        For ``MODE_EAX``, ``MODE_GCM`` and ``MODE_SIV`` there are no
        restrictions on its length (recommended: **16** bytes).

        For ``MODE_CCM``, its length must be in the range **[7..13]**.
        Bear in mind that with CCM there is a trade-off between nonce
        length and maximum message size. Recommendation: **11** bytes.

        For ``MODE_OCB``, its length must be in the range **[1..15]**
        (recommended: **15**).

        For ``MODE_CTR``, its length must be in the range **[0..15]**
        (recommended: **8**).

        For ``MODE_SIV``, the nonce is optional, if it is not specified,
        then no nonce is being used, which renders the encryption
        deterministic.

        If not provided, for modes other than ``MODE_SIV``, a random
        byte string of the recommended length is used (you must then
        read its value with the :attr:`nonce` attribute).

      segment_size (integer):
        (Only ``MODE_CFB``).The number of **bits** the plaintext and ciphertext
        are segmented in. It must be a multiple of 8.
        If not specified, it will be assumed to be 8.

      mac_len (integer):
        (Only ``MODE_EAX``, ``MODE_GCM``, ``MODE_OCB``, ``MODE_CCM``)
        Length of the authentication tag, in bytes.

        It must be even and in the range **[4..16]**.
        The recommended value (and the default, if not specified) is **16**.

      msg_len (integer):
        (Only ``MODE_CCM``). Length of the message to (de)cipher.
        If not specified, ``encrypt`` must be called with the entire message.
        Similarly, ``decrypt`` can only be called once.

      assoc_len (integer):
        (Only ``MODE_CCM``). Length of the associated data.
        If not specified, all associated data is buffered internally,
        which may represent a problem for very large messages.

      initial_value (integer or bytes/bytearray/memoryview):
        (Only ``MODE_CTR``).
        The initial value for the counter. If not present, the cipher will
        start counting from 0. The value is incremented by one for each block.
        The counter number is encoded in big endian mode.

      counter (object):
        (Only ``MODE_CTR``).
        Instance of ``Crypto.Util.Counter``, which allows full customization
        of the counter block. This parameter is incompatible to both ``nonce``
        and ``initial_value``.

      use_aesni: (boolean):
        Use the AES instructions of the CPU (default: use if available):
        AES-NI on x86, or the Cryptography Extensions on 64-bit ARM.

      threads (integer):
        (Only ``MODE_CTR``, ``MODE_GCM`` and ``MODE_OCB``).
        The maximum number of threads used to encrypt or decrypt long data
        (default: 1, no extra threads).
        Use 0 for as many threads as the CPU cores available
        to this process.
        Each thread processes at least 1 MiB of a single call to
        ``encrypt()`` or ``decrypt()``: with shorter inputs,
        fewer threads are used, or none at all.
        For best results, do not exceed the number of physical cores.
        The output does not depend on the number of threads.

    Returns:
        :ref:`A cipher object <aes_modes>`: Its class depends on ``mode``,
        see the table in :ref:`aes_modes`.
    """

    kwargs["add_aes_modes"] = True
    return _create_cipher(sys.modules[__name__], key, mode, *args, **kwargs)


#: Size of a data block (in bytes)
block_size: int = 16
#: The valid sizes of a key (in bytes)
key_size: tuple[int, int, int] = (16, 24, 32)