:- module(xchacha20_poly1305, [
    test_xchacha20_poly1305/0,
    xchacha20_poly1305_encrypt/6,
    xchacha20_poly1305_decrypt/6
]).
:- use_module(chacha20_poly1305).
:- use_module(chacha20_prim).


xchacha20_poly1305_encrypt(Aad, Key, Nonce, Text, Ciphertext, Tag) :-
    assertion(is_u8s(Aad)),
    assertion(is_xchacha20_key(Key)),
    assertion(is_xchacha20_nonce(Nonce)),
    assertion(is_u8s(Text)),

    split_nonce(N0, N1, Nonce),
    hchacha20(Key, N0, Subkey),
    chacha20_poly1305_encrypt(Aad, Subkey, N1, 0, Text, Ciphertext, Tag),
    
    assertion(is_u8s(Ciphertext)),
    assertion(is_poly1305_tag(Tag)).

xchacha20_poly1305_decrypt(Aad, Key, Nonce, Text, Ciphertext, Tag) :-
    assertion(is_u8s(Aad)),
    assertion(is_xchacha20_key(Key)),
    assertion(is_xchacha20_nonce(Nonce)),
    assertion(is_u8s(Ciphertext)),
    assertion(is_u8s(Tag)),

    split_nonce(N0, N1, Nonce),
    hchacha20(Key, N0, Subkey),
    chacha20_poly1305_decrypt(Aad, Subkey, N1, 0, Text, Ciphertext, Tag),

    assertion(is_u8s(Text)).

split_nonce(N0, N1, Nonce) :-
    length(N0, 16),
    length(N1, 8),
    append(N0, N1, Nonce).

hchacha20(Key, Nonce, Hash) :-
    assertion(is_chacha20_key(Key)),
    assertion(is_u8s(Nonce)), length(Nonce, 16),
    u8s_to_u32s_le(Key, K),
    u8s_to_u32s_le(Nonce, N),
    hchacha20_impl(K, N, H),
    u8s_to_u32s_le(Hash, H),
    assertion(is_hchacha20_hash(Hash)).

hchacha20_impl(Key, Nonce, Hash) :-
    hchacha20_block(Key, Nonce, [
        H0, H1, H2, H3, 
        _, _, _, _, 
        _, _, _, _, 
        H4, H5, H6, H7
    ]),
    Hash = [H0, H1, H2, H3, H4, H5, H6, H7].

hchacha20_block(Key, Nonce, Output) :-
    Key = [K0, K1, K2, K3, K4, K5, K6, K7],
    Nonce = [N0, N1, N2, N3],
    Input = [
        0x61707865, 0x3320646e, 0x79622d32, 0x6b206574,
        K0, K1, K2, K3,
        K4, K5, K6, K7,
        N0, N1, N2, N3
    ],
    chacha20_prim_u32(Input, Output).

test_xchacha20_poly1305 :-
    test_hchacha20,
    test_e2e.

test_hchacha20 :-
    hchacha20_block(
        [
            0x03020100, 0x07060504, 0x0b0a0908, 0x0f0e0d0c,
            0x13121110, 0x17161514, 0x1b1a1918, 0x1f1e1d1c
        ],
        [0x09000000, 0x4a000000, 0x00000000, 0x27594131],
        Output
    ),
    assertion(
        Output = [
            0x423b4182, 0xfe7bb227, 0x50420ed3, 0x737d878a,
            0x0aa76448, 0x7954cdf3, 0x846acd37, 0x7b3c58ad,
            0x77e35583, 0x83e77c12, 0xe0076a2d, 0xbc6cd0e5,
            0xd5e4f9a0, 0x53a8748a, 0x13c42ec1, 0xdcecd326
        ]
    ).

test_e2e :-
    hex_bytes("50515253c0c1c2c3c4c5c6c7", Aad),
    hex_bytes("808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f", Key),
    hex_bytes("404142434445464748494a4b4c4d4e4f5051525354555657", Nonce),
    string_bytes(
        "Ladies and Gentlemen of the class of \x27\99: If I could offer you only one tip for the future, sunscreen would be it.",
        Text, utf8
    ),
    hex_bytes("bd6d179d3e83d43b9576579493c0e939572a1700252bfaccbed2902c21396cbb731c7f1b0b4aa6440bf3a82f4eda7e39ae64c6708c54c216cb96b72e1213b4522f8c9ba40db5d945b11b69b982c1bb9e3f3fac2bc369488f76b2383565d3fff921f9664c97637da9768812f615c68b13b52e", ExpectedCiphertext),
    hex_bytes("c0875924c1c7987947deafd8780acf49", ExpectedTag),

    xchacha20_poly1305_encrypt(Aad, Key, Nonce, Text, Ciphertext, Tag),

    assertion(Ciphertext = ExpectedCiphertext),
    assertion(Tag = ExpectedTag).