/*-
 * Copyright 2021 John-Mark Gurney.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS 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 ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 *
 */

#include <comms.h>
#include <strobe_rng_init.h>

static const size_t MAC_LEN = 8;
static const size_t CHALLENGE_LEN = 16;
static const uint8_t shared_domain[] = "com.funkthat.lora.irrigation.shared.v0.0.1";
static const uint8_t ecdhe_domain[] = "com.funkthat.lora.irrigation.ecdhe.v0.0.1";
static const uint8_t reqreset[] = "reqreset";
static const uint8_t confirm[] = "confirm";

static int comms_pktbuf_equal(struct pktbuf a, struct pktbuf b);

/* returns 1 if equal, 0 if not equal */
static int
comms_pktbuf_equal(struct pktbuf a, struct pktbuf b)
{

	if (a.pktlen != b.pktlen)
		return 0;

	return memcmp(a.pkt, b.pkt, a.pktlen) == 0;
}

size_t
_comms_state_size()
{

	return sizeof(struct comms_state);
}

size_t
_strobe_state_size()
{

	return sizeof(strobe_s);
}

int
comms_init(struct comms_state *cs, process_msgfunc_t pmf,
    struct pktbuf *shared, struct pktbuf *resp_key, struct pktbuf *init_pubkey)
{
	unsigned char buf[EC_PUBLIC_BYTES * 2];
	strobe_s *crypt;

	*cs = (struct comms_state){
		.cs_procmsg = pmf,
	};

	crypt = &cs->cs_start.cs_crypto;

	if (shared != NULL) {
		strobe_init(crypt, shared_domain, sizeof shared_domain - 1);

		strobe_key(crypt, SYM_KEY, shared->pkt, shared->pktlen);

		cs->cs_start.cs_state = COMMS_WAIT_REQUEST_SHARED;
	} else if (resp_key != NULL && init_pubkey != NULL) {
		strobe_init(crypt, ecdhe_domain, sizeof ecdhe_domain - 1);

		if (resp_key->pktlen != EC_PRIVATE_BYTES ||
		    init_pubkey->pktlen != EC_PUBLIC_BYTES)
			return 0;

		/* store our private key */
		memcpy(cs->cs_respkey, resp_key->pkt, resp_key->pktlen);
		x25519_base(cs->cs_resppubkey, resp_key->pkt, 1);

		/* store the public key to associate with */
		memcpy(cs->cs_initpubkey, init_pubkey->pkt, init_pubkey->pktlen);

		/* public keys */
		memcpy(&buf[0], cs->cs_initpubkey, EC_PUBLIC_BYTES);
		memcpy(&buf[EC_PUBLIC_BYTES], cs->cs_resppubkey, EC_PUBLIC_BYTES);

		strobe_key(crypt, SYM_KEY, buf, EC_PUBLIC_BYTES * 2);

		cs->cs_start.cs_state = COMMS_WAIT_REQUEST_ECDHE;
	} else {
		return 0;
	}

	/* copy starting state over to initial state */
	cs->cs_active = cs->cs_start;
	cs->cs_pending = cs->cs_active;

	return 1;
}

#define CONFIRMED_STR_BASE	"confirmed"
#define CONFIRMED_STR		((const uint8_t *)CONFIRMED_STR_BASE)
#define CONFIRMED_STR_LEN	(sizeof(CONFIRMED_STR_BASE) - 1)

static void
_comms_process_session(struct comms_state *cs, struct comms_session *sess, struct pktbuf pbin, struct pktbuf *pbout)
{
	strobe_s tmp;
	uint8_t buf[64] = {};
	struct pktbuf pbmsg, pbrep;
	ssize_t cnt, ret, msglen;

	/* save the state incase the message is bad */
	tmp = sess->cs_crypto;

	strobe_attach_buffer(&sess->cs_crypto, pbin.pkt, pbin.pktlen);

	/* if the packet is too short, ignore */
	if (pbin.pktlen < MAC_LEN)
		goto badmsg;

	cnt = strobe_get(&sess->cs_crypto, APP_CIPHERTEXT, buf, pbin.pktlen -
	    MAC_LEN);
	msglen = cnt;

	cnt = strobe_get(&sess->cs_crypto, MAC, pbin.pkt +
	    (pbin.pktlen - MAC_LEN), MAC_LEN);

	/* MAC check failed */
	if (cnt == -1) {
badmsg:
		/* restore the previous state */
		sess->cs_crypto = tmp;
		pbout->pktlen = 0;
		return;
	}

	/*
	 * if we have arrived here, MAC has been verified, and buf now
	 * contains the data to operate upon.
	 */

	/* attach the buffer for output */
	strobe_attach_buffer(&sess->cs_crypto, pbout->pkt, pbout->pktlen);

	ret = 0;
	switch (sess->cs_state) {
	case COMMS_WAIT_REQUEST_SHARED:
		if (msglen != 24 || memcmp(reqreset, &buf[16],
		    sizeof reqreset - 1) != 0)
			goto badmsg;

		bare_strobe_randomize(buf, CHALLENGE_LEN);
		ret = strobe_put(&sess->cs_crypto, APP_CIPHERTEXT, buf,
		    CHALLENGE_LEN);
		ret += strobe_put(&sess->cs_crypto, MAC, NULL, MAC_LEN);

		strobe_operate(&sess->cs_crypto, RATCHET, NULL, 32);

		sess->cs_state = COMMS_WAIT_CONFIRM;
		break;

	case COMMS_WAIT_REQUEST_ECDHE:
	{
		unsigned char ephkey[EC_PRIVATE_BYTES];
		unsigned char ephpubkey[EC_PUBLIC_BYTES];
		unsigned char keybuf[EC_PUBLIC_BYTES * 2];

		if (msglen != (EC_PUBLIC_BYTES + 8) || memcmp(reqreset, &buf[EC_PUBLIC_BYTES],
		    sizeof reqreset - 1) != 0)
			goto badmsg;

		/* DH(s, re) */
		memcpy(&keybuf[0], cs->cs_resppubkey, EC_PUBLIC_BYTES);
		if (x25519(&keybuf[0], cs->cs_respkey, buf, 1) != 0)
			goto badmsg;

		/* DH(s, rs) */
		memcpy(&keybuf[EC_PUBLIC_BYTES], cs->cs_resppubkey, EC_PUBLIC_BYTES);
		if (x25519(&keybuf[EC_PUBLIC_BYTES], cs->cs_respkey, cs->cs_initpubkey, 1) != 0)
			goto badmsg;

		/* key(DH() + DH()) */
		strobe_key(&sess->cs_crypto, SYM_KEY, keybuf, EC_PUBLIC_BYTES * 2);

		/* generate ephemeral */
		bare_strobe_randomize(ephkey, EC_PRIVATE_BYTES);
		if (x25519_base(ephpubkey, ephkey, 1) != 0)
			goto badmsg;

		ret = strobe_put(&sess->cs_crypto, APP_CIPHERTEXT, ephpubkey,
		    EC_PUBLIC_BYTES);
		ret += strobe_put(&sess->cs_crypto, MAC, NULL, MAC_LEN);

		/* DH(e, re) */
		memcpy(&keybuf[0], ephpubkey, EC_PUBLIC_BYTES);
		if (x25519(&keybuf[0], ephkey, buf, 1) != 0)
			goto badmsg;

		/* DH(e, rs) */
		memcpy(&keybuf[EC_PUBLIC_BYTES], ephpubkey, EC_PUBLIC_BYTES);
		if (x25519(&keybuf[EC_PUBLIC_BYTES], ephkey, cs->cs_initpubkey, 1) != 0)
			goto badmsg;

		/* key(DH() + DH()) */
		strobe_key(&sess->cs_crypto, SYM_KEY, keybuf, EC_PUBLIC_BYTES * 2);

		sess->cs_state = COMMS_WAIT_CONFIRM;
		break;
	}

	case COMMS_WAIT_CONFIRM:
		if (msglen != 7 || memcmp(confirm, buf,
		    sizeof confirm - 1) != 0)
			goto badmsg;

		ret = strobe_put(&sess->cs_crypto, APP_CIPHERTEXT, CONFIRMED_STR,
		    CONFIRMED_STR_LEN);
		ret += strobe_put(&sess->cs_crypto, MAC, NULL, MAC_LEN);
		sess->cs_state = COMMS_PROCESS_MSGS;
		break;

	case COMMS_PROCESS_MSGS: {
		uint8_t repbuf[pbout->pktlen - MAC_LEN];

		memset(repbuf, '\x00', sizeof repbuf);

		pbmsg.pkt = buf;
		pbmsg.pktlen = msglen;

		pbrep.pkt = repbuf;
		pbrep.pktlen = sizeof repbuf;

		cs->cs_procmsg(pbmsg, &pbrep);

		ret = strobe_put(&sess->cs_crypto, APP_CIPHERTEXT, repbuf,
		    pbrep.pktlen);
		ret += strobe_put(&sess->cs_crypto, MAC, NULL, MAC_LEN);

		break;
		}
	}

	/* set the output buffer length */
	pbout->pktlen = ret;

}

/*
 * encrypted data to be processed is passed in via pbin.
 *
 * The pktbuf pointed to by pbout contains the buffer that a [encrypted]
 * response will be written to.  The length needs to be updated, where 0
 * means no reply.
 */
void
comms_process(struct comms_state *cs, struct pktbuf pbin, struct pktbuf *pbout)
{
	struct pktbuf pbouttmp;

	/* if the current msg matches the previous */
	if (comms_pktbuf_equal(pbin, cs->cs_prevmsg)) {
		/* send the previous response */
		pbout->pktlen = cs->cs_prevmsgresp.pktlen;
		memcpy(pbout->pkt, cs->cs_prevmsgresp.pkt, pbout->pktlen);
		return;
	}

	/* try to use the active session */
	pbouttmp = *pbout;
	_comms_process_session(cs, &cs->cs_active, pbin, &pbouttmp);

	if (pbouttmp.pktlen != 0) {
retmsg:
		/* we accepted a new message store it */
		*pbout = pbouttmp;

		/* store the req */
		cs->cs_prevmsg.pkt = cs->cs_prevmsgbuf;
		cs->cs_prevmsg.pktlen = pbin.pktlen;
		memcpy(cs->cs_prevmsg.pkt, pbin.pkt, pbin.pktlen);

		/* store the response */
		cs->cs_prevmsgresp.pkt = cs->cs_prevmsgrespbuf;
		cs->cs_prevmsgresp.pktlen = pbout->pktlen;
		memcpy(cs->cs_prevmsgresp.pkt, pbout->pkt, pbout->pktlen);
	} else {
		/* active session didn't work, try cs_pending */

		pbouttmp = *pbout;
		_comms_process_session(cs, &cs->cs_pending, pbin, &pbouttmp);

		if (cs->cs_pending.cs_state == COMMS_PROCESS_MSGS) {
			/* new active state */
			cs->cs_active = cs->cs_pending;
			cs->cs_pending = cs->cs_start;
			goto retmsg;
		}

		/* pending session advanced (likely to _WAIT_CONFIRM) */
		if (pbouttmp.pktlen > 0) {
			*pbout = pbouttmp;
			return;
		}

		/* pending session didn't work, maybe new */
		struct comms_session tmpsess;

		tmpsess = cs->cs_start;

		pbouttmp = *pbout;
		_comms_process_session(cs, &tmpsess, pbin, &pbouttmp);
		if (tmpsess.cs_state == COMMS_WAIT_CONFIRM) {
			/* new request for session */
			cs->cs_pending = tmpsess;
			*pbout = pbouttmp;
		} else {
			/* no packet to reply with */
			pbout->pktlen = 0;
		}
	}
}