:- module(chacha20, [chacha20_block/5, chacha20_cipher/5, test_chacha20/0]).
:- use_module(chacha20_prim).
:- use_module(types).

chacha20_block(Key, Nonce, Counter, State, Mixed) :-
    assertion(is_chacha20_key(Key)),
    assertion(is_chacha20_nonce(Nonce)),
    assertion(is_chacha20_counter(Counter)),
    u8s_to_u32s_le(Key, [K0, K1, K2, K3, K4, K5, K6, K7]),
    expand_counter(Counter, [C0, C1]),
    u8s_to_u32s_le(Nonce, [N0, N1]),
    BlockStart = [
        % specified here: https://datatracker.ietf.org/doc/html/rfc7539
        0x61707865, 0x3320646e, 0x79622d32, 0x6b206574,
        K0, K1, K2, K3,
        K4, K5, K6, K7,
        C0, C1, N0, N1
    ],
    % format("\n"),
    % format("start block: "), print_hexlist(BlockStart), format("\n"),
    chacha20_prim_u32(BlockStart, State),
    % format("state block: "), print_hexlist(State), format("\n"),
    maplist(mix_term, BlockStart, State, Mixed).
    % format("mixed block: "), print_hexlist(Mixed), format("\n"),
    % format("\n").

mix_term(X, Y, Z) :- Z is (X + Y) /\ 0xFFFFFFFF.

chacha20_cipher(Key, Nonce, Counter, Text, Result) :-
    assertion(is_chacha20_key(Key)),
    assertion(is_chacha20_nonce(Nonce)),
    assertion(is_chacha20_counter(Counter)),
    assertion(is_u8s(Text)),
    chacha20_cipher_impl(Key, Nonce, Counter, Text, Result).

chacha20_cipher_impl(_, _, _, [],[]).
chacha20_cipher_impl(Key, Nonce, Counter, Text, Result) :-
    append_chunk(64, TextPrefix, TextRemaining, Text),
    compute_keystream(Key, Nonce, Counter, Keystream),
    mix_keystream(Keystream, TextPrefix, ResultPrefix),
    assertion(length(Keystream, 64)),
    chacha20_cipher_impl(Key, Nonce, Counter + 1, TextRemaining, ResultRemaining),
    append(ResultPrefix, ResultRemaining, Result).

compute_keystream(Key, Nonce, Counter, Keystream) :- 
    chacha20_block(Key, Nonce, Counter, _, Mixed),
    u8s_to_u32s_le(Keystream, Mixed).

expand_counter(Counter, [C1, C2]) :- 
    C1 is Counter mod (1 << 32),
    C2 is Counter div (1 << 32).

mix_keystream(_, [], []).
mix_keystream([K|Key], [T|Text], [O|Out]) :-
    O is K xor T,
    mix_keystream(Key, Text, Out).
mix_keystream([], [_|_], _) :- throw(ran_out_of_keystream).

test_chacha20 :- 
    test_compute_keystream,
    test_mix_keystream,
    test_e2e_short,
    test_e2e_long.

test_compute_keystream :-
    % test vectors from here: https://datatracker.ietf.org/doc/html/rfc7539#appendix-A.1
    ZeroKey = [
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
    ],
    OneKey = [
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1
    ],
    HighKey = [
        0, 0xff, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0   , 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
    ],
    ZeroNonce = [0, 0, 0, 0, 0, 0, 0, 0],
    TwoNonce = [0, 0, 0, 0, 0, 0, 0, 2],
    forall(
        member(example(K, N, Ctr) = Ks, [
            example(ZeroKey, ZeroNonce, 0) = [0x76, 0xb8, 0xe0, 0xad],
            example(ZeroKey, ZeroNonce, 1) = [0x9f, 0x07, 0xe7, 0xbe],
            example(OneKey, ZeroNonce, 1) = [0x3a, 0xeb, 0x52, 0x24],
            example(HighKey, ZeroNonce, 2) = [0x72, 0xd5, 0x4d, 0xfb],
            example(ZeroKey, TwoNonce, 0) = [0xc2, 0xc6, 0x4d, 0x37]
        ]),
        assertion((
            compute_keystream(K, N, Ctr, Ks2), 
            prefix(Ks, Ks2)
        ))
    ).

test_mix_keystream :-
    forall(
        member(mix(Ks, Tx) = Result, [
            mix([0, 0, 0], []) = [],
            mix([0, 0, 0], [0, 0, 0]) = [0, 0, 0],
            mix([1, 2, 3], [8, 8, 8]) = [9, 10, 11],
            mix([1, 2, 3], [9, 10, 11]) = [8, 8, 8]
        ]),
        assertion( mix_keystream(Ks, Tx, Result) )
    ).

test_e2e_short :-
    length(PlaintextBytes, 64), maplist(=(0), PlaintextBytes),
    Ciphertext="76b8e0ada0f13d90405d6ae55386bd28bdd219b8a08ded1aa836efcc8b770dc7da41597c5157488d7724e03fb8d84a376a43b8f41518a11cc387b669b2ee6586",
    hex_bytes(Ciphertext, CiphertextBytes),
    zeros(32, Key), zeros(8, Nonce),
    chacha20_cipher(
        Key, Nonce, 0,
        PlaintextBytes,
        CiphertextCandidateBytes
    ),
    assertion(CiphertextBytes = CiphertextCandidateBytes),
    chacha20_cipher(
        Key, Nonce, 0,
        CiphertextBytes,
        PlaintextCandidateBytes
    ),
    assertion(PlaintextBytes = PlaintextCandidateBytes).

test_e2e_long :-
    Plaintext = "Any submission to the IETF intended by the Contributor for publication as all or part of an IETF Internet-Draft or RFC and any statement made within the context of an IETF activity is considered an \"IETF Contribution\". Such statements include oral statements in IETF sessions, as well as written and electronic communications made at any time or place, which are addressed to",
    Ciphertext = "a3fbf07df3fa2fde4f376ca23e82737041605d9f4f4f57bd8cff2c1d4b7955ec2a97948bd3722915c8f3d337f7d370050e9e96d647b7c39f56e031ca5eb6250d4042e02785ececfa4b4bb5e8ead0440e20b6e8db09d881a7c6132f420e52795042bdfa7773d8a9051447b3291ce1411c680465552aa6c405b7764d5e87bea85ad00f8449ed8f72d0d662ab052691ca66424bc86d2df80ea41f43abf937d3259dc4b2d0dfb48a6c9139ddd7f76966e928e635553ba76c5c879d7b35d49eb2e62b0871cdac638939e25e8a1e0ef9d5280fa8ca328b351c3c765989cbcf3daa8b6ccc3aaf9f3979c92b3720fc88dc95ed84a1be059c6499b9fda236e7e818b04b0bc39c1e876b193bfe5569753f88128cc08aaa9b63d1a16f80ef2554d7189c411f5869ca52c5b83fa36ff216b9c1d30062bebcfd2dc5bce0911934fda79a86f6e698ced759c3ff9b6477338f3da4f9cd8514ea9982ccafb341b2384dd902f3d1ab7ac61dd29c6f21ba5b862f3730e37cfdc4fd806c22f221",
    string_bytes(Plaintext, PlaintextBytes, utf8),
    hex_bytes(Ciphertext, CiphertextBytes),
    Key = [
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1
    ],
    Nonce = [0, 0, 0, 0, 0, 0, 0, 2],
    chacha20_cipher(
        Key, Nonce, 1,
        PlaintextBytes,
        CiphertextCandidateBytes
    ),
    assertion(CiphertextBytes = CiphertextCandidateBytes),
    chacha20_cipher(
        Key, Nonce, 1,
        CiphertextBytes,
        PlaintextCandidateBytes
    ),
    assertion(PlaintextBytes = PlaintextCandidateBytes).