/*
 * Copyright (C) 2022 - 2026, Stephan Mueller <smueller@chronox.de>
 *
 * License: see LICENSE file in root directory
 *
 * THIS SOFTWARE IS PROVIDED ``AS IS'' AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE, ALL OF
 * WHICH ARE HEREBY DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT
 * OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 * USE OF THIS SOFTWARE, EVEN IF NOT ADVISED OF THE POSSIBILITY OF SUCH
 * DAMAGE.
 */

#include "compare.h"
#include "ext_headers_internal.h"
#include "lc_ascon_lightweight.h"
#include "test_helper_common.h"
#include "visibility.h"

static int ascon_tester_one(const uint8_t *pt, size_t ptlen,
			    const uint8_t *nonce, size_t noncelen,
			    const uint8_t *aad, size_t aadlen,
			    const uint8_t *key, size_t keylen,
			    const uint8_t *exp_ct, const uint8_t *exp_tag,
			    size_t exp_tag_len)
{
	struct lc_aead_ctx *al_heap = NULL;
	ssize_t ret;
	int ret_checked = 0;
	uint8_t out_enc[32];
	uint8_t out_dec[32];
	uint8_t tag[16];
	LC_AL_CTX_ON_STACK(al);

	if (ptlen > sizeof(out_enc))
		return -EINVAL;

	if (exp_tag_len > sizeof(tag))
		return -EINVAL;

	/* One shot encryption with pt ptr != ct ptr */
	if (lc_aead_setkey(al, key, keylen, nonce, noncelen))
		return 1;
	lc_aead_encrypt(al, pt, out_enc, ptlen, aad, aadlen, tag, exp_tag_len);
	lc_aead_zero(al);

	ret_checked +=
		lc_compare(out_enc, exp_ct, ptlen,
			   "Ascon lightweight crypt: Encryption, ciphertext");
	ret_checked += lc_compare(tag, exp_tag, exp_tag_len,
				  "Ascon lightweight crypt: Encryption, tag");

	//bin2print(out_enc, ptlen, stderr, "out_enc");
	//bin2print(tag, exp_tag_len, stderr, "tag");

	/* One shot encryption with pt ptr == ct ptr */
	if (lc_al_alloc(&al_heap))
		return 1;

	if (lc_aead_setkey(al_heap, key, keylen, nonce, noncelen))
		return 1;

	if (ptlen)
		memcpy(out_enc, pt, ptlen);
	lc_aead_encrypt(al_heap, out_enc, out_enc, ptlen, aad, aadlen, tag,
			exp_tag_len);
	lc_aead_zero_free(al_heap);

	ret_checked += lc_compare(out_enc, exp_ct, ptlen,
				  "Ascon lightweight: Encryption, ciphertext");
	ret_checked += lc_compare(tag, exp_tag, exp_tag_len,
				  "Ascon lightweight: Encryption, tag");

	/* One shot decryption with pt ptr != ct ptr */
	if (lc_aead_setkey(al, key, keylen, nonce, noncelen))
		return -EFAULT;
	ret = lc_aead_decrypt(al, out_enc, out_dec, ptlen, aad, aadlen, tag,
			      exp_tag_len);

	//bin2print(out_dec, ptlen, stderr, "out_enc");
	lc_aead_zero(al);
	if (ret < 0)
		ret_checked += 1;

	ret_checked +=
		lc_compare(out_dec, pt, ptlen,
			   "Ascon lightweight crypt: Decryption, plaintext");

	/* Check authentication error */
	if (lc_aead_setkey(al, key, keylen, nonce, noncelen))
		return -EFAULT;

	if (ptlen) {
		out_enc[0] = (uint8_t)((out_enc[0] + 1) & 0xff);
		ret = lc_aead_decrypt(al, out_enc, out_dec, ptlen, aad, aadlen,
				      tag, exp_tag_len);
		if (ret != -EBADMSG)
			ret_checked += 1;
	}

	lc_aead_zero(al);

	return ret_checked;
}

static int ascon_tester_128(void)
{
	/*
	 * Vector 1089 from genkat_crypto_aead_asconaead128_ref generated by
	 * code https://github.com/ascon/ascon-c
	 */
	static const uint8_t pt[] = { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06,
				      0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D,
				      0x0E, 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14,
				      0x15, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x1B,
				      0x1C, 0x1D, 0x1E, 0x1F };
	static const uint8_t aad[] = { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06,
				       0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D,
				       0x0E, 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14,
				       0x15, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x1B,
				       0x1C, 0x1D, 0x1E, 0x1F };
	static const uint8_t key[] = { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05,
				       0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B,
				       0x0C, 0x0D, 0x0E, 0x0F };
	static const uint8_t nonce[] = { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05,
					 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B,
					 0x0C, 0x0D, 0x0E, 0x0F };
	static const uint8_t exp_ct[] = { 0x4C, 0x08, 0x6D, 0x27, 0xA3, 0xB5,
					  0x1A, 0x23, 0x33, 0xCF, 0xC7, 0xF2,
					  0x21, 0x72, 0xA9, 0xBC, 0xAD, 0x88,
					  0xB8, 0xD4, 0xD7, 0x7E, 0x50, 0x62,
					  0x2D, 0x78, 0x83, 0x45, 0xFA, 0x7B,
					  0xEE, 0x44 };
	static const uint8_t exp_tag[] = { 0x68, 0x91, 0x5D, 0x3F, 0x94, 0x22,
					   0x28, 0x9F, 0x23, 0x49, 0xD6, 0xA3,
					   0xB4, 0x16, 0x03, 0x97 };
	int ret;

	printf("Ascon lightweight 128 crypt ctx len %u, state len %u\n",
	       (unsigned int)LC_AL_CTX_SIZE, (unsigned int)LC_AL_STATE_SIZE);

	ret = ascon_tester_one(pt, sizeof(pt), nonce, sizeof(nonce), aad,
			       sizeof(aad), key, sizeof(key), exp_ct, exp_tag,
			       sizeof(exp_tag));

	return ret;
}

static int ascon_tester_128_non_aligned(void)
{
	static const uint8_t pt[] = { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06,
				      0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D,
				      0x0E, 0x0F, 0x10, 0x11, 0x12, 0x13, 0x14,
				      0x15, 0x16, 0x17, 0x18, 0x19, 0x1A, 0x1B,
				      0x1C, 0x1D, 0x1E };
	static const uint8_t key[] = { 0x00, 0x01, 0x02, 0x03, 0x04, 0x05,
				       0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B,
				       0x0C, 0x0D, 0x0E, 0x0F };
	static const uint8_t exp_ct[] = { 0x14, 0x1B, 0x8B, 0x25, 0xE5, 0x9E,
					  0x0D, 0x01, 0xB1, 0x14, 0x22, 0xD9,
					  0x40, 0x48, 0xC8, 0x04, 0x5E, 0xE7,
					  0x74, 0x4D, 0xD0, 0x04, 0x0B, 0xF0,
					  0x8B, 0x6C, 0xED, 0x0E, 0x4F, 0x25,
					  0x5B };
	static const uint8_t exp_tag[] = { 0x19, 0x6B, 0x88, 0x31, 0x37, 0xE0,
					   0xD1, 0xCA, 0xF3, 0x3A, 0xFE, 0xBB,
					   0x4B, 0x9E, 0x72, 0xC4 };

	return ascon_tester_one(pt, sizeof(pt), key, sizeof(key), pt,
				sizeof(pt), key, sizeof(key), exp_ct, exp_tag,
				sizeof(exp_tag));
}

static int ascon_tester_128_null_aad(void)
{
	/* From NIST ACVP */
	static const uint8_t pt[] = { 0xA6, 0xA6 };
	static const uint8_t nonce[] = { 0xE0, 0x92, 0xE3, 0x90, 0xF8, 0xD1,
					 0x2E, 0xB9, 0x15, 0x52, 0x0B, 0x2C,
					 0xE2, 0x93, 0xE6, 0xBA };
	static const uint8_t key[] = { 0x53, 0x1C, 0x5B, 0x05, 0xE6, 0x16,
				       0x9D, 0x2B, 0xEE, 0x5D, 0xF8, 0x0C,
				       0xCC, 0x8A, 0x2A, 0xAF };
	static const uint8_t exp_ct[] = { 0xc9, 0xb8 };
	static const uint8_t exp_tag[] = { 0x93, 0x43, 0x81, 0x7a, 0x83, 0x2b,
					   0xe0, 0xec, 0xd8, 0x8e, 0x3f, 0xfc,
					   0x0d, 0x81, 0x51, 0x1b };

	return ascon_tester_one(pt, sizeof(pt), nonce, sizeof(nonce), NULL, 0,
				key, sizeof(key), exp_ct, exp_tag,
				sizeof(exp_tag));
}

static int ascon_tester_128_null_auth_only(void)
{
	/* From NIST ACVP */
	static const uint8_t nonce[] = { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
					 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
					 0x00, 0x00, 0x00, 0x00 };
	static const uint8_t key[] = { 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
				       0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
				       0x00, 0x00, 0x00, 0x00 };
	static const uint8_t exp_tag[] = { 0xf1, 0xe8, 0xda, 0x91, 0xf9, 0xfb,
					   0xf0, 0x90, 0xca, 0xf3, 0xe3, 0x7c,
					   0xca, 0xd9, 0x1b, 0xae };

	return ascon_tester_one(NULL, 0, nonce, sizeof(nonce), NULL, 0, key,
				sizeof(key), NULL, exp_tag, sizeof(exp_tag));
}

static int ascon_tester_128_null_pt(void)
{
	/* From NIST ACVP */
	static const uint8_t aad[] = { 0x50, 0x9C };
	static const uint8_t nonce[] = { 0xE0, 0xE5, 0xE0, 0x06, 0x67, 0xDB,
					 0xD1, 0x6B, 0x40, 0x92, 0x6D, 0x1F,
					 0xAD, 0xED, 0xE7, 0x79 };
	static const uint8_t key[] = { 0x5E, 0xF6, 0x37, 0x72, 0x00, 0xDC,
				       0xAB, 0x0C, 0x05, 0xD3, 0x3C, 0x14,
				       0xE7, 0x4D, 0x83, 0x3F };
	static const uint8_t exp_tag[] = { 0x4c, 0xd9, 0x61, 0x3a, 0x76, 0x79,
					   0x82, 0x5a, 0xc5, 0x79, 0x4a, 0x83,
					   0xd3, 0x45, 0xe4, 0x29 };

	return ascon_tester_one(NULL, 0, nonce, sizeof(nonce), aad, sizeof(aad),
				key, sizeof(key), NULL, exp_tag,
				sizeof(exp_tag));
}

static int ascon_tester_128_null_pt_aad(void)
{
	/* From NIST ACVP */
	static const uint8_t nonce[] = { 0x0F, 0x0B, 0xB5, 0x97, 0xD2, 0x87,
					 0x97, 0x53, 0x44, 0x53, 0x49, 0x87,
					 0x00, 0x96, 0xE2, 0x9B };
	static const uint8_t key[] = { 0x24, 0x66, 0x17, 0xC0, 0xF9, 0xFD,
				       0x0A, 0x1D, 0x3F, 0x04, 0x5D, 0x92,
				       0x04, 0xBF, 0x71, 0xE6 };
	static const uint8_t exp_tag[] = { 0x0e, 0xfa, 0x20, 0x72, 0x16, 0x28,
					   0x9f, 0x85, 0xe2, 0xaf, 0xed, 0x13,
					   0xd8, 0x6c, 0x21, 0x81 };

	return ascon_tester_one(NULL, 0, nonce, sizeof(nonce), NULL, 0, key,
				sizeof(key), NULL, exp_tag, sizeof(exp_tag));
}

LC_TEST_FUNC(int, main, int argc, char *argv[])
{
	int ret = 0;
	(void)argc;
	(void)argv;

	ret += ascon_tester_128_null_auth_only();
	ret += ascon_tester_128();
	ret += ascon_tester_128_non_aligned();
	ret += ascon_tester_128_null_aad();
	ret += ascon_tester_128_null_pt();
	ret += ascon_tester_128_null_pt_aad();

	ret = test_validate_status(ret, lc_aead_alg_status(lc_ascon_aead), 1);
	ret += test_print_status();

	return ret;
}
