/* implementation of sponge based crypto by eeeee
 * Public Domain
 *
 * Cryptographic sponge functions:
 * https://keccak.team/files/CSF-0.1.pdf
 * https://web.archive.org/web/20250620013520/https://keccak.team/files/CSF-0.1.pdf
 *
 * Sponge functions require permutation function. I use gimli:
 * https://gimli.cr.yp.to/
 * https://gimli.cr.yp.to/gimli-20170627.pdf
 * https://web.archive.org/web/20250421113716/https://gimli.cr.yp.to/gimli-20170627.pdf
 */
#ifndef SPONGE_BOB_H
#define SPONGE_BOB_H

#include <arpa/inet.h>
#include <string.h>
#include "shorttypes.h"

#define ROTL32(x, n) (((x)<<(n)) | ((x)>>(32-(n))))

static void gimli(void *s){
	u32 round, column;
	u32 x, y, z;
	u32 *state = s;

	for (round = 24; round != 0; --round){
		for (column = 0; column < 4; ++column){
			x = ROTL32(state[    column], 24);
			y = ROTL32(state[4 + column],  9);
			z =        state[8 + column];

			state[8 + column] = x ^ (z << 1) ^ ((y&z) << 2);
			state[4 + column] = y ^ x        ^ ((x|z) << 1);
			state[    column] = z ^ y        ^ ((x&y) << 3);
		}

		if ((round & 3) == 0) { // small swap: pattern s...s...s...
			x = state[0];
			state[0] = state[1];
			state[1] = x;
			x = state[2];
			state[2] = state[3];
			state[3] = x;
		}

		if ((round & 3) == 2) { // big swap: ..S...S...S. etc.
			x = state[0];
			state[0] = state[2];
			state[2] = x;
			x = state[1];
			state[1] = state[3];
			state[3] = x;
		}

		if ((round & 3) == 0) { // add constant: c...c...c... etc.
			state[0] ^= (0x9e377900 | round);
		}
	}

	return;
}

/* gimli, endian-agnostic version */
static void gimli_ea(void *s){
	u32 *state = s;

	state[ 0] = ntohl(state[ 0]);
	state[ 1] = ntohl(state[ 1]);
	state[ 2] = ntohl(state[ 2]);
	state[ 3] = ntohl(state[ 3]);
	state[ 4] = ntohl(state[ 4]);
	state[ 5] = ntohl(state[ 5]);
	state[ 6] = ntohl(state[ 6]);
	state[ 7] = ntohl(state[ 7]);
	state[ 8] = ntohl(state[ 8]);
	state[ 9] = ntohl(state[ 9]);
	state[10] = ntohl(state[10]);
	state[11] = ntohl(state[11]);

	gimli(s);

	state[ 0] = htonl(state[ 0]);
	state[ 1] = htonl(state[ 1]);
	state[ 2] = htonl(state[ 2]);
	state[ 3] = htonl(state[ 3]);
	state[ 4] = htonl(state[ 4]);
	state[ 5] = htonl(state[ 5]);
	state[ 6] = htonl(state[ 6]);
	state[ 7] = htonl(state[ 7]);
	state[ 8] = htonl(state[ 8]);
	state[ 9] = htonl(state[ 9]);
	state[10] = htonl(state[10]);
	state[11] = htonl(state[11]);

	return;
}

#ifndef FPERMUTE
#define FPERMUTE gimli_ea
#endif

#define MEMXOR16(dst, src) \
	((u32*)dst)[0] ^= ((u32*)src)[0]; \
	((u32*)dst)[1] ^= ((u32*)src)[1]; \
	((u32*)dst)[2] ^= ((u32*)src)[2]; \
	((u32*)dst)[3] ^= ((u32*)src)[3]

#define MEMXOR32(dst, src) \
	((u32*)dst)[0] ^= ((u32*)src)[0]; \
	((u32*)dst)[1] ^= ((u32*)src)[1]; \
	((u32*)dst)[2] ^= ((u32*)src)[2]; \
	((u32*)dst)[3] ^= ((u32*)src)[3]; \
	((u32*)dst)[4] ^= ((u32*)src)[4]; \
	((u32*)dst)[5] ^= ((u32*)src)[5]; \
	((u32*)dst)[6] ^= ((u32*)src)[6]; \
	((u32*)dst)[7] ^= ((u32*)src)[7]


/* reseedable prng
 * state[12]: input buffer index
 * state[13]: output buffer index
 */
static void duplex257_prng_feed(u32 state[14], u32 food){
	state[13] = 0;
	state[state[12]] ^= food;
	state[12]++;
	if (state[12] == 8){
		((u8*)state)[32] ^= 0x80;
		state[12] = 0;
		FPERMUTE(state);
	}
	return;
}

static u32 duplex257_prng_rand(u32 state[14]){
	u32 res;
	if (state[12]){
		((u8*)state)[state[12]] ^= 0x80;
		FPERMUTE(state);
		state[12] = 0;
		state[13] = 1;
		return state[0];
	}

	res = state[state[13]];
	state[13]++;
	if (state[13] == 8){
		state[13] = 0;
		((u8*)state)[0] ^= 0x80;
		FPERMUTE(state);
	}

	return res;
}

static void duplex257_prng_rand16(u32 state[14], u8 dst[16]){
	if (state[12]){
		((u8*)state)[state[12]] ^= 0x80;
		FPERMUTE(state);
		state[12] = 0;
		state[13] = 4;
		memcpy(dst, state, 16);
		return;
	}

	if (state[13] == 0){
		state[13] = 4;
		memcpy(dst, state, 16);
	} else if (state[13] < 5){
		state[13] = 0;
		memcpy(dst, &state[4], 16);
		((u8*)state)[0] ^= 0x80;
		FPERMUTE(state);
	} else {
		((u8*)state)[0] ^= 0x80;
		FPERMUTE(state);
		state[13] = 4;
		memcpy(dst, state, 16);
	}
	return;
}

/* duplex257 AE */
static void duplex257_ae_encrypt
    (const u8 key[32], const u8 nonce[16], u8 tag[16], u8 *src_dst, u32 sz){
	static u8 state[48];
	u32 i, n;

	/* duplex_mute(key) */
	memcpy(state, key, 32);
	state[32] = 0x80;
	FPERMUTE(state);

	/* duplex_mute(nonce) */
	MEMXOR16(state, nonce);
	state[16] ^= 0x80;
	FPERMUTE(state);

	/* 32 bytes steps */
	n = sz >> 5;
	while (n){
		/* duplex(src_dst), encrypt */
		MEMXOR32(src_dst, state);
		memcpy(state, src_dst, 32);
		state[32] ^= 0x80;
		FPERMUTE(state);
		src_dst += 32;
		n--;
	}

	/* 1 byte steps */
	n = sz & 0x1Fu;
	/* duplex(src_dst), encrypt */
	for (i = 0; i < n; i++) src_dst[i] ^= state[i];
	memcpy(state, src_dst, n);
	state[n] ^= 0x80;
	FPERMUTE(state);

	memcpy(tag, state, 16);
	memset(state, 0, sizeof(state));
	return;
}

static void duplex257_ae_decrypt
    (const u8 key[32], const u8 nonce[16], u8 tag[16], u8 *src_dst, u32 sz){
	static u8 state[48];
	u32 i, n;

	/* duplex_mute(key) */
	memcpy(state, key, 32);
	state[32] = 0x80;
	FPERMUTE(state);

	/* duplex_mute(nonce) */
	MEMXOR16(state, nonce);
	state[16] ^= 0x80;
	FPERMUTE(state);

	/* 32 bytes steps */
	n = sz >> 5;
	while (n){
		/* duplex(src_dst), encrypt */
		MEMXOR32(src_dst, state);
		MEMXOR32(state, src_dst);
		state[32] ^= 0x80;
		FPERMUTE(state);
		src_dst += 32;
		n--;
	}

	/* 1 byte steps */
	n = sz & 0x1Fu;
	/* duplex(src_dst), encrypt */
	for (i = 0; i < n; i++){
		src_dst[i] ^= state[i];
		state[i] ^= src_dst[i];
	}
	state[n] ^= 0x80;
	FPERMUTE(state);

	memcpy(tag, state, 16);
	memset(state, 0, sizeof(state));
	return;
}

/* sponge */
static void sponge256_absorb(u8 state[48], const u8 *src, u32 sz){
	u32 n, i;

	n = sz >> 5;
	while (n){ /* 32 byte steps */
		MEMXOR32(state, src);
		FPERMUTE(state);
		src += 32;
		n--;
	}

	/* 1 byte steps */
	n = sz & 0x1Fu;
	for (i=0; i < n; i++) state[i] ^= src[i];
	state[i] ^= 0x80;
	FPERMUTE(state);
	return;
}

/* sponge hash */
static void sponge256_hash32(u8 dst[32], u8 *src, u32 sz){
	static u8 state[48];
	sponge256_absorb(state, src, sz);
	memcpy(dst, state, 32);
	memset(state, 0, sizeof(state));
	return;
}

#endif /* SPONGE_BOB_H */