module AsconCipher where

import Ascon

// 4. Authenticated Encryption Schema: Ascon-AEAD128

Ascon_AEAD128 : {a, p} (fin a, fin p) => [128] -> [128] -> [a] -> [p] -> [p + 128]
Ascon_AEAD128 K N A P = C # T
  where
    // key and nonce as two 64-bit integers
    [K0,K1] = bitsToWords K
    [N0,N1] = bitsToWords N

    S0 = Ascon_p`{12} [Ascon_AEAD128_IV, K0, K1, N0, N1]
       ^ [0, 0, 0, K0, K1]
    
    SA = AddAD S0 A
  
    SCs = zipWith XorBlock (take ([SA] # map Ascon_p`{8} SCs)) (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_bytes : {a, p} (fin a, fin p) => [16][8] -> [16][8] -> [a][8] -> [p][8] -> [p + 16][8]
Ascon_AEAD128_bytes K N A P =
  bitsToWords (Ascon_AEAD128 (wordsToBits K) (wordsToBits N) (wordsToBits A) (wordsToBits P))

AddAD : {a} (fin a) => State -> [a] -> State
AddAD S A
  | a == 0 => DomainSep S
  | a > 0  => DomainSep (foldl AbsorbBlock S (toBlocks A))

XorBlock : State -> [128] -> State
XorBlock [s0, s1, s2, s3, s4] (xhi # xlo) = [s0 ^ xlo, s1 ^ xhi, s2, s3, s4]

AbsorbBlock : State -> [128] -> State
AbsorbBlock S X = Ascon_p`{8} (XorBlock S X)

DomainSep : State -> State
DomainSep [s0, s1, s2, s3, s4] = [s0, s1, s2, s3, s4 ^ 0b1 # 0]

ExtractC : State -> [128]
ExtractC [s0, s1, _, _, _] = wordsToBits [s0, s1]

ExtractT : State -> [128]
ExtractT [_, _, _, s3, s4] = wordsToBits [s3, s4]

Ascon_AEAD128_IV : [64]
Ascon_AEAD128_IV = 0x00001000808c0001