module aes

import encoding.hex

// Test vectors from "The Galois/Counter Mode of Operation (GCM)",
// and NIST SP 800-38D worked examples.
fn test_aes128_gcm_vector() {
    key := hex.decode('feffe9928665731c6d6a8f9467308308')!
    iv := hex.decode('cafebabefacedbaddecaf888')!
    plaintext :=
        hex.decode('d9313225f88406e5a55909c5aff5269a86a7a9531534f7da2e4c303d8a318a721c3c0c95956809532fcf0e2449a6b525b16aedf5aa0de657ba637b39')!
    ad := hex.decode('feedfacedeadbeeffeedfacedeadbeefabaddad2')!
    want_ct :=
        hex.decode('42831ec2217774244b7221b784d0d49ce3aa212f2c02a4e035c17e2329aca12e21d514b25466931c7d8f6a5aac84aa051ba30b396a0aac973d58e091')!
    want_tag := hex.decode('5bc94fbc3221a5db94fae95ae7121a47')!

    mut g := new_aes_gcm(key)!
    out := g.encrypt(plaintext, iv, ad)!
    ct := out[..out.len - 16]
    tag := out[out.len - 16..]
    assert ct == want_ct
    assert tag == want_tag

    dec := g.decrypt(out, iv, ad)!
    assert dec == plaintext
}

fn test_aes256_gcm_vector() {
    key := hex.decode('feffe9928665731c6d6a8f9467308308feffe9928665731c6d6a8f9467308308')!
    iv := hex.decode('cafebabefacedbaddecaf888')!
    plaintext :=
        hex.decode('d9313225f88406e5a55909c5aff5269a86a7a9531534f7da2e4c303d8a318a721c3c0c95956809532fcf0e2449a6b525b16aedf5aa0de657ba637b39')!
    ad := hex.decode('feedfacedeadbeeffeedfacedeadbeefabaddad2')!
    want_ct :=
        hex.decode('522dc1f099567d07f47f37a32a84427d643a8cdcbfe5c0c97598a2bd2555d1aa8cb08e48590dbb3da7b08b1056828838c5f61e6393ba7a0abcc9f662')!
    want_tag := hex.decode('76fc6ece0f4e1768cddf8853bb2d551b')!

    mut g := new_aes_gcm(key)!
    out := g.encrypt(plaintext, iv, ad)!
    assert out[..out.len - 16] == want_ct
    assert out[out.len - 16..] == want_tag
    assert g.decrypt(out, iv, ad)! == plaintext
}

fn test_gcm_tamper_detected() {
    key := hex.decode('00000000000000000000000000000000')!
    iv := hex.decode('000000000000000000000000')!
    mut g := new_aes_gcm(key)!
    out := g.encrypt('hello'.bytes(), iv, []u8{})!
    mut bad := out.clone()
    bad[0] ^= 0xff
    if _ := g.decrypt(bad, iv, []u8{}) {
        assert false, 'tampered ciphertext should fail authentication'
    }
}

// The test materials taken and adapted from CAVP Testing: Block Cipher Modes.
// Its only take the samples, not all tests was performed, some of them was not supported.
// See https://csrc.nist.gov/CSRC/media/Projects/Cryptographic-Algorithm-Validation-Program/documents/mac/gcmtestvectors.zip
fn test_avcp_1() ! {
    // [Keylen = 128]
    // [IVlen = 96]
    // [PTlen = 0]
    // [AADlen = 0]
    // [Taglen = 128]

    // Count = 0
    key := hex.decode('11754cd72aec309bf52f7687212e8957')!
    iv := hex.decode('3c819d9a9bed087615030b65')!
    pt := hex.decode('')!
    aad := hex.decode('')!
    ct := hex.decode('')!
    tag := hex.decode('250327c674aaf477aef2675748cf6971')!

    mut g := new_aes_gcm(key)!
    out := g.encrypt(pt, iv, aad)!

    // ciphertext only contains the tag
    assert out == tag
}

fn test_avcp_2() ! {
    // [Keylen = 128]
    // [IVlen = 96]
    // [PTlen = 128]
    // [AADlen = 128]
    // [Taglen = 120] // 15 bytes
    //
    // Count = 0
    key := hex.decode('89c54b0d3bc3c397d5039058c220685f')!
    iv := hex.decode('bc7f45c00868758d62d4bb4d')!
    pt := hex.decode('582670b0baf5540a3775b6615605bd05')!
    aad := hex.decode('48d16cda0337105a50e2ed76fd18e114')!
    ct := hex.decode('fc2d4c4eee2209ddbba6663c02765e69')!
    tag := hex.decode('55e783b00156f5da0446e2970b877f')!

    mut g := new_aes_gcm(key)!
    out := g.encrypt(pt, iv, aad)!
    assert out[0..pt.len] == ct
    // the tag len was truncated to 15-bytes
    assert out[pt.len..pt.len + 15].hex() == tag.hex()

    mut g2 := new_aes_gcm(key)!
    ptx := g2.decrypt(out, iv, aad)!
    assert ptx == pt
}

struct ItemTest {
    count int
    key   string
    iv    string
    pt    string
    aad   string
    ct    string
    tag   string
}

fn test_avcp_3() ! {
    // [Keylen = 128]
    // [IVlen = 96]
    // [PTlen = 128]
    // [AADlen = 160]
    // [Taglen = 128]
    tag_len := 16

    for t in tests_data {
        key := hex.decode(t.key)!
        iv := hex.decode(t.iv)!
        pt := hex.decode(t.pt)!
        aad := hex.decode(t.aad)!
        ct := hex.decode(t.ct)!
        tag := hex.decode(t.tag)!

        mut g := new_aes_gcm(key)!
        out := g.encrypt(pt, iv, aad)!
        assert out[0..pt.len] == ct
        assert out[pt.len..pt.len + tag_len] == tag
    }
}

const tests_data = [
    ItemTest{
        count: 0
        key:   'd4a22488f8dd1d5c6c19a7d6ca17964c'
        iv:    'f3d5837f22ac1a0425e0d1d5'
        pt:    '7b43016a16896497fb457be6d2a54122'
        aad:   'f1c5d424b83f96c6ad8cb28ca0d20e475e023b5a'
        ct:    'c2bd67eef5e95cac27e3b06e3031d0a8'
        tag:   'f23eacf9d1cdf8737726c58648826e9c'
    },
    ItemTest{
        count: 1
        key:   'e8899345e4d89b76f7695ddf2a24bb3c'
        iv:    '9dfaeb5d73372ceb06ca7bbe'
        pt:    'c2807e403e9babf645268c92bc9d1de6'
        aad:   'fed0b45a9a7b07c6da5474907f5890e317e74a42'
        ct:    '8e44bf07454255aa9e36eb34cdfd0036'
        tag:   '2f501e5249aa595a53e1985e90346a22'
    },
    ItemTest{
        count: 2
        key:   'c1629d6320b9da80a23c81be53f0ef57'
        iv:    'b8615f6ffa30668947556cd8'
        pt:    '65771ab52532c9cdfcb3a9eb7b8193df'
        aad:   '5f2955e4301852a70684f978f89e7a61531f0861'
        ct:    'c2a72d693181c819f69b42b52088d3a2'
        tag:   'cadaee305d8bb6d70259a6503280d99a'
    },
    ItemTest{
        count: 3
        key:   '196ed78281bb7543d60e68cca2aaa941'
        iv:    '6e7d2c8f135715532a075c50'
        pt:    '15b42e7ea21a8ad5dcd7a9bba0253d44'
        aad:   'd6fc98c632d2e2641041ff7384d92a8358ae9abe'
        ct:    '06e5cc81c2d022cb2b5de5a881c62d09'
        tag:   '28e8cad3346ce583d5eebaa796e50974'
    },
    ItemTest{
        count: 4
        key:   '55fe8a1bdc6806ed2f4a84891db943a0'
        iv:    'af4d0ba0a90f1e713d71ae94'
        pt:    '81315972f0b1aeaa005363e9eca09d7a'
        aad:   '677cd4e6c0a67913085dba4cc2a778b894e174ad'
        ct:    'c47bcb27c5a8d9beb19fee38b90861b7'
        tag:   'e061ee4868edf2d969e875b8685ca8a9'
    },
    ItemTest{
        count: 5
        key:   '6d86a855508657f804091be2290a17e0'
        iv:    '65dce18a4461afd83f1480f5'
        pt:    '0423bd1c8aea943637c7c3b0ca61d54b'
        aad:   'e0ef8f0e1f442a2c090568d2af336ec59f57c896'
        ct:    '53505d449369c9bcd8a138740ea6602e'
        tag:   '86f928b4532825af9cac3820234afe73'
    },
    ItemTest{
        count: 6
        key:   '66bd7b5dfd0aaaed8bb8890eee9b9c9a'
        iv:    '6e92bf7e8fd0fb932451fdf2'
        pt:    '8005865c8794b79612447f5ef33397d0'
        aad:   '60459c681bda631ece1aacca4a7b1b369c56d2bb'
        ct:    '83b99253de05625aa8e68490bb368bb9'
        tag:   '65d444b02a23e854a85423217562d07f'
    },
]