/** Ascon-Based Lightweight Cryptography
 *
 * Author: Eric Mertens <emertens@gmail.com>
 * License: ISC
 *
 * Key algorithms:
 * - AEAD128_encrypt/decrypt: Authenticated encryption
 * - Hash256: Cryptographic hash function
 * - XOF128: Extendable output function
 * - CXOF128: Customizable extendable output function
 *
 * Reference:
 *
 * Meltem Sönmez Turan, Kerry A. McKay, Donghoon Chang, Jinkeon Kang,
 * John Kelsey (2025) Ascon-Based Lightweight Cryptography Standards
 * for Constrained Devices. (National Institute of Standards and Technology,
 * Gaithersburg, MD), NIST Special Publication (SP) NIST SP 800-232.
 * https://doi.org/10.6028/NIST.SP.800-232
 */
module Ascon where

private
  // 2.1. Auxiliary Functions

  /** Parse function.
   *
   * The parse(𝑋, 𝑟) function parses the input bitstring 𝑋 into a sequence
   * of blocks 𝑋₀, 𝑋₁, …, 𝑋͠ℓ, where 𝓁 ← ⌊|𝑋|/𝑟⌋ (i.e., 𝑋 ← 𝑋₀ ∥ 𝑋₁ ∥ … ∥ 𝑋͠ℓ).
   * The 𝑋ᵢ blocks for 0 ≤ i ≤ 𝓁 − 1 each have a bit length 𝑟, whereas
   * 0 ≤ |𝑋͠ℓ| ≤ 𝑟 − 1 (see Algorithm 1). When |𝑋| mod 𝑟 = 0, the final
   * block is empty (i.e., |𝑋͠ℓ| = 0).
   */
  parse : {r, m} (fin m, fin r, r >= 1) => [m] -> ([m / r][r], [m % r])
  parse (M # Ml) = (split M, Ml)

  /** Padding rule.
   *
   * The function pad(𝑋, 𝑟) appends the bit 1 to the bitstring 𝑋, followed
   * by the bitstring 0ʲ, where 𝑗 is equal to (−|𝑋|−1) mod 𝑟. The length of
   * the output bitstring is a multiple of 𝑟 (see Algorithm 2). For examples
   * of padding when representing the data as 64-bit unsigned integers, see
   * Appendix A.2.
   */
  pad : {r, m} (m < r, fin r) => [m] -> [r]
  pad M = M # 0b1 # 0

  /**
   * Combination of parse and pad that splits a bitstring into a sequence
   * of integers using Cryptol's native big-endian representation.
   */
  toBlocks : {r, m} (r >= 1, fin r, fin m) => [m] -> [m / r + 1][r]
  toBlocks M = bitsToWords (M # 0b1 # 0)

  // 3. Ascon Permutations

  /** Predicate on valid round counts. */
  type constraint ValidRnd rnd = rnd <= 16

  /** Core permutation function parameterized by a round count. */
  Ascon_p : {rnd} (ValidRnd rnd) => State -> State
  Ascon_p S = foldl round S (drop`{back=rnd} Const)

  /**
  * Single round of the Ascon-p permutation parameterized by the round
  * constant.
  */
  round : State -> [64] -> State
  round S ci = pL (pS (pC S ci))

  // 3.1. Internal State

  /**
  * The permutations operate on the 320-bit state 𝑆, which is represented
  * as five 64-bit words denoted as 𝑆ᵢ for 0 ≤ 𝑖 ≤ 4:
  *
  *     𝑆 = 𝑆₀ ‖ 𝑆₁ ‖ 𝑆₂ ‖ 𝑆₃ ‖ 𝑆₄
  *
  * Let 𝑠ᵢⱼ represent the 𝑗th bit of 𝑆ᵢ, 0 ≤ 𝑗 < 64. In this specification
  * of the Ascon permutation, each state word represents a 64-bit unsigned
  * integer, where the least significant bit is the rightmost bit. Details
  * on other representations of the state can be found in Appendix A.
  */
  type State = [5][64]

  // 3.2. Constant-Addition Layer pC

  // The constant cᵢ of round 𝑖 of the Ascon permutation Ascon-p[𝑟𝑛𝑑]
  // (instantiated with 𝑟𝑛𝑑 rounds) for 𝑟𝑛𝑑 ≤ 16 and 0 ≤ 𝑖 < 𝑟𝑛𝑑 − 1
  // is defined as
  //
  //     cᵢ = const[16 − 𝑟𝑛𝑑 + 𝑖]
  //
  // where const[0],…,const[15] are defined in Table 5.

  /**
  * The constant-addition layer 𝑃𝑐 adds a 64-bit round constant cᵢ to 𝑆₂
  * in round 𝑖, for i ≥ 0, 𝑆₂ = 𝑆₂ ⊕ cᵢ.
  */
  pC : State -> [64] -> State
  pC [S0, S1, S2, S3, S4] ci = [S0, S1, S2 ^ ci, S3, S4]

  /**
  * Table 5. The constants constᵢ to derive round constants of the Ascon
  * permutations
  */
  Const : [16][64]
  Const =
      [ 0x000000000000003c
      , 0x000000000000002d
      , 0x000000000000001e
      , 0x000000000000000f
      , 0x00000000000000f0
      , 0x00000000000000e1
      , 0x00000000000000d2
      , 0x00000000000000c3
      , 0x00000000000000b4
      , 0x00000000000000a5
      , 0x0000000000000096
      , 0x0000000000000087
      , 0x0000000000000078
      , 0x0000000000000069
      , 0x000000000000005a
      , 0x000000000000004b
      ]

  // 3.3. Substitution Layer pS

  /**
   * The substitution layer 𝑝𝑆 updates the state S with 64 parallel
   * applications of the 5-bit substitution box SBOX
   */
  pS : State -> State
  pS S = transpose (map SBox (transpose S))

  /** 5-bit substitution used in pS. */
  SBox : [5] -> [5]
  SBox i = SBoxTable@i

  /** Table 6. Lookup table for SBox */
  SBoxTable : [32][5]
  SBoxTable =
      map drop
      [ 0x04, 0x0b, 0x1f, 0x14, 0x1a, 0x15, 0x09, 0x02
      , 0x1b, 0x05, 0x08, 0x12, 0x1d, 0x03, 0x06, 0x1c
      , 0x1e, 0x13, 0x07, 0x0e, 0x00, 0x0d, 0x11, 0x18
      , 0x10, 0x0c, 0x01, 0x19, 0x16, 0x0a, 0x0f, 0x17
      ]

  property
    SBoxCorrect x0 x1 x2 x3 x4 = [y0, y1, y2, y3, y4] == SBox [x0, x1, x2, x3, x4]
      where
        y0 = x4&&x1 ^ x3 ^ x2&&x1 ^ x2 ^ x1&&x0 ^ x1 ^ x0
        y1 = x4 ^ x3&&x2 ^ x3&&x1 ^ x3 ^ x2&&x1 ^ x2 ^ x1 ^ x0
        y2 = x4&&x3 ^ x4 ^ x2 ^ x1 ^ 1
        y3 = x4&&x0 ^ x4 ^ x3&&x0 ^ x3 ^ x2 ^ x1 ^ x0
        y4 = x4&&x1 ^ x4 ^ x3 ^ x1&&x0 ^ x1

  // 3.4. Linear Diffusion Layer pL

  /** The linear diffusion layer 𝑝𝐿 provides diffusion within each 64-bit word 𝑆𝑖. */
  pL : State -> State
  pL [S0, S1, S2, S3, S4] =
      [ sigma S0 19 28
      , sigma S1 61 39
      , sigma S2  1  6
      , sigma S3 10 17
      , sigma S4  7 41
      ]
      where
          sigma : [64] -> [6] -> [6] -> [64]
          sigma x i j = x ^ x>>>i ^ x>>>j

  /** Concatenate a sequence of integers into a little-endian bitstream. */
  wordsToBits : {w, n} (fin w) => [n][w] -> [w*n]
  wordsToBits M = join (map reverse M)

  /** Split a little-endian bitstream in to a sequence of integers. */
  bitsToWords : {w, n} (fin w) => [w*n] -> [n][w]
  bitsToWords M = map reverse (split M)

// 4. Authenticated Encryption Schema: Ascon-AEAD128

/** Ascon-AEAD128 encryption algorithm on bitstreams
 *
 * Type parameters:
 * - a: Bit-length of associated data
 * - p: Bit-length of plaintext
 *
 * Parameters:
 * - K: Key
 * - N: Nonce
 * - A: Associated data
 * - P: Plaintext
 *
 * Returns:
 * - Authenticated ciphertext
 */
AEAD128_encrypt : {a, p} (fin a, fin p) => [128] -> [128] -> [a] -> [p] -> [p + 128]
AEAD128_encrypt K N A P = C # T
  where
    [K0, K1] = bitsToWords K
    [N0, N1] = bitsToWords N

    // 𝑆 ← 𝐼𝑉 ‖ 𝐾 ‖ 𝑁              (15)
    // 𝑆 ← Ascon-p[12](𝑆)          (16)
    // 𝑆 ← 𝑆 ⊕ (0¹⁹² ‖ K)          (17)
    S0 = Ascon_p`{12} [AEAD128_IV, K0, K1, N0, N1] ^ [0, 0, 0, K0, K1]

    SA = AddAD S0 A

    // 𝑃₀, 𝑃₁, …, 𝑃ₙ₋₁, 𝑃͠ₙ ← parse(𝑃,128)       (23)
    // For each 𝑃ᵢ, 0 ≤ 𝑖 ≤ 𝑛−1
    //     𝑆[0:127] ← 𝑆[0:127] ⊕ 𝑃ᵢ             (24)
    //     𝐶ᵢ ← 𝑆[0:127]                        (25)
    //     𝑆 ← Ascon-p[8](𝑆)                    (26)
    SCs = [XorBlock s p | s <- [SA] # map Ascon_p`{8} SCs | p <- toBlocks P]
    C = take (join (map ExtractC SCs))

    ST = Ascon_p`{12} (last SCs ^ [0, 0, K0, K1, 0])
    T = ExtractT ST ^ K

/** Ascon-AEAD128 decryption algorithm on bitstreams.
 *
 * Type parameters:
 * - a: Bit-length of associated data
 * - p: Bit-length of plaintext
 *
 * Parameters:
 * - K: Key
 * - N: Nonce
 * - A: Associated data
 * - C: Ciphertext
 *
 * Returns:
 * - Some plaintext on authentication success
 * - None on authentication failure
 */
AEAD128_decrypt : {a, p} (fin a, fin p) => [128] -> [128] -> [a] -> [p + 128] -> Option [p]
AEAD128_decrypt K N A (C # T) = if T == T' then Some P else None
  where
    // key and nonce as two 64-bit integers
    [K0, K1] = bitsToWords K
    [N0, N1] = bitsToWords N

    S0 = Ascon_p`{12} [AEAD128_IV, K0, K1, N0, N1] ^ [0, 0, 0, K0, K1]
    SA = AddAD S0 A

    (Cs, Cl1) = parse C
    SCs = [SA] # [Ascon_p`{8} (AssignC s c) | s <- SCs | c <- Cs]

    Mask # Cl2 = join (map ExtractC SCs)
    P = Mask ^ C

    Cl = Cl1 # (0b1#0 ^ Cl2) // xor toggles the padding bit
    ST1 = AssignC (last SCs) Cl
    ST2 = Ascon_p`{12} (ST1 ^ [0, 0, K0, K1, 0])
    T' = ExtractT ST2 ^ K

/** Ascon-AEAD128 encryption algorithm on bytes.
 *
 * Type parameters:
 * - a: Byte-length of associated data
 * - p: Byte-length of plaintext
 *
 * Parameters:
 * - K: Key
 * - N: Nonce
 * - A: Associated data
 * - P: Plaintext
 *
 * Returns:
 * - Authenticated ciphertext
 */
AEAD128_encrypt_bytes : {a, p} (fin a, fin p) => [16][8] -> [16][8] -> [a][8] -> [p][8] -> [p + 16][8]
AEAD128_encrypt_bytes K N A P =
  bitsToWords (AEAD128_encrypt (wordsToBits K) (wordsToBits N) (wordsToBits A) (wordsToBits P))

/** Ascon-AEAD128 decryption algorithm on bytes.
 *
 * Type parameters:
 * - a: Byte-length of associated data
 * - p: Byte-length of plaintext
 *
 * Parameters:
 * - K: Key
 * - N: Nonce
 * - A: Associated data
 * - C: Ciphertext
 *
 * Returns:
 * - Some plaintext on authentication success
 * - None on authentication failure
 */
AEAD128_decrypt_bytes : {a, p} (fin a, fin p) => [16][8] -> [16][8] -> [a][8] -> [p + 16][8] -> Option ([p][8])
AEAD128_decrypt_bytes K N A C =
  case AEAD128_decrypt (wordsToBits K) (wordsToBits N) (wordsToBits A) (wordsToBits C) of
    None -> None
    Some p -> Some (bitsToWords p)

private
  /** Absorb all of the associated data into the permutation state. */
  AddAD : {a} (fin a) => State -> [a] -> State
  AddAD S A

      /* If the AD is non-empty (i.e., |𝐴| > 0):
       *    𝐴₀, 𝐴₁, …, 𝐴ₘ₋₁, 𝐴͠ₘ ← parse(𝐴,128)   (18)
       *    𝐴ₘ ← pad(𝐴͠ₘ,128)                     (19)
       */
      | a > 0  => DomainSep (foldl AbsorbADBlock S (toBlocks A))

      /* If the AD is empty (i.e., |𝐴| = 0): Only the final step
       * described in (22) is applied.
       */
      | a == 0 => DomainSep S

  /* Absorb a block in input into the state.
   *     𝑆[0:127] ← 𝑆[0:127] ⊕ 𝐴ᵢ      (20)
   *     𝑆 ← Ascon-p[8](𝑆)             (21)
   **/
  AbsorbADBlock : State -> [128] -> State
  AbsorbADBlock S Ai = Ascon_p`{8} (XorBlock S Ai)

  /** Toggle the domain separation bit indicating end of associated data.
   *    𝑆 ← 𝑆 ⊕ (0³¹⁹ ‖ 1)             (22)
   */
  DomainSep : State -> State
  DomainSep [s0, s1, s2, s3, s4] = [s0, s1, s2, s3, s4 ^ 0b1#0]

  /** Add a single block of data into the permutation state. */
  XorBlock : State -> [128] -> State
  XorBlock [s0, s1, s2, s3, s4] (xhi # xlo) = [s0 ^ xlo, s1 ^ xhi, s2, s3, s4]

  /** Extracts the first two words of permutation state as a bitstream. */
  ExtractC : State -> [128]
  ExtractC [s0, s1, _, _, _] = wordsToBits [s0, s1]

  /** Assigns a bitstream into the first two words of permutation state. */
  AssignC : State -> [128] -> State
  AssignC [_, _, s2, s3, s4] C = bitsToWords C # [s2, s3, s4]

  /** Extracts the last two words of permutation state as a bitstream. */
  ExtractT : State -> [128]
  ExtractT [_, _, _, s3, s4] = wordsToBits [s3, s4]

  /** Ascon-AEAD128 initialization vector */
  AEAD128_IV : [64]
  AEAD128_IV = 0x00001000808c0001

// 5. Hash and eXtendable-Output Functions (XOFs)

private
  /** Hash function construction
   *
   * Parameters:
   * - Initialization vector
   * - List of blocks as integers
   *
   * Returns:
   * - Infinite message digest bitstream to be truncated by caller
   */
  hashBlocks : {n} (fin n) => [64] -> [n][64] -> [inf]
  hashBlocks IV Ms = wordsToBits [head S | S <- iterate Ascon_p`{12} Sn]
      where
      S0 = Ascon_p`{12} [IV, 0, 0, 0, 0]
      Sn = foldl AbsorbBlock64 S0 Ms

  AbsorbBlock64 : State -> [64] -> State
  AbsorbBlock64 [s0, s1, s2, s3, s4] X = Ascon_p`{12} [X ^ s0, s1, s2, s3, s4]

// 5.1. Specification of Ascon-Hash256

/** Ascon-Hash256 implementation on bitstreams.
 *
 * Type parameters:
 * - m: Bit-length of message
 *
 * Parameters:
 * - M: Message
 *
 * Returns:
 * - Message digest
 */
Hash256 : {m} (fin m) => [m] -> [256]
Hash256 M = take (hashBlocks Hash256_IV (toBlocks M))

/** Ascon-Hash256 implementation on bytes.
 *
 * Type parameters:
 * - m: Byte-length of message
 *
 * Parameters:
 * - M: Message
 *
 * Returns:
 * - Message digest
 */
Hash256_bytes : {m} (fin m) => [m][8] -> [32][8]
Hash256_bytes M = bitsToWords (Hash256 (wordsToBits M))

/** Ascon-Hash256 initialization vector */
Hash256_IV : [64]
private Hash256_IV = 0x0000080100cc0002

// 5.2. Specification of Ascon-XOF128

/** Ascon-XOF256 implementation on bitstreams.
 *
 * Type parameters:
 * - r: Bit-length of digest
 * - m: Bit-length of message
 *
 * Parameters:
 * - M: Message
 *
 * Returns:
 * - Variable-length message digest
 */
XOF128 : {r, m} (fin m, fin r) => [m] -> [r]
XOF128 M = take (hashBlocks XOF128_IV (toBlocks M))

/** Ascon-XOF256 implementation on bytes.
 *
 * Type parameters:
 * - r: Byte-length of digest
 * - m: Byte-length of message
 *
 * Parameters:
 * - M: Message
 *
 * Returns:
 * - Variable-length message digest
 */
XOF128_bytes : {r, n} (fin n, fin r) => [n][8] -> [r][8]
XOF128_bytes M = bitsToWords (XOF128 (wordsToBits M))

/** Ascon-XOF128 initialization vector */
XOF128_IV : [64]
private XOF128_IV = 0x0000080000cc0003

// 5.3. Specification of Ascon-CXOF128

/** Ascon-CXOF256 implementation on bitstreams.
 *
 * Type parameters:
 * - r: Bit-length of digest
 * - c: Bit-length of customization string
 * - m: Bit-length of message
 *
 * Parameters:
 * - Z: User-defined customization string
 * - M: Message
 *
 * Returns:
 * - Variable-length message digest
 */
CXOF128 : {r, c, m} (fin m, fin r, fin c, 64 >= width c) => [c] -> [m] -> [r]
CXOF128 Z M = take (hashBlocks CXOF128_IV Ms)
  where
    Ms = [`c]
       # toBlocks Z
       # toBlocks M

/** Ascon-CXOF256 implementation on bytes.
 *
 * Type parameters:
 * - r: Byte-length of digest
 * - c: Byte-length of customization string
 * - m: Byte-length of message
 *
 * Parameters:
 * - Z: User-defined customization string
 * - M: Message
 *
 * Returns:
 * - Variable-length message digest
 */
CXOF128_bytes : {r, z, m} (fin m, fin r, 61 >= width z) => [z][8] -> [m][8] -> [r][8]
CXOF128_bytes Z M = bitsToWords (CXOF128 (wordsToBits Z) (wordsToBits M))

/** Ascon-CXOF128 initialization vector */
CXOF128_IV : [64]
private CXOF128_IV = 0x0000080000cc0004