#define _POSIX_C_SOURCE 200809L

#include <arpa/inet.h>

#include "tlse.h"

#ifndef htonll
#define htonll(x) ((1==htonl(1)) ? (x) : ((uint64_t)htonl((x) & 0xFFFFFFFF) << 32) | htonl((x) >> 32))
#endif

#ifndef ntohll
#define ntohll(x) ((1==ntohl(1)) ? (x) : ((uint64_t)ntohl((x) & 0xFFFFFFFF) << 32) | ntohl((x) >> 32))
#endif

static int crypto_encrypt(struct TLSContext *context, unsigned char *buf,
		unsigned char *ct, unsigned int len) {
	if (context->crypto.created == 1) {
		return cbc_encrypt(buf, ct, len,
				   &context->crypto.ctx_local.aes_local);
	}

	memset(ct, 0, len);
	return TLS_GENERIC_ERROR;
}

static int packet_encrypt(struct TLSPacket *packet) {
	int header_size = 5;
	int block_size = TLS_AES_BLOCK_SIZE;
	int mac_size = 0;
	unsigned int length = 0;
	unsigned char padding = 0;
	//unsigned int pt_length = packet->len - header_size;
	struct TLSContext *context = packet->context;

	mac_size = tls_mac_length(packet->context);
	length = packet->len - header_size + TLS_AES_IV_LENGTH + mac_size;

	padding = block_size - length % block_size;
	length += padding;
	unsigned char *buf = malloc(length);
	if (buf) {
		unsigned char *ct = malloc(length + header_size);
		if (ct) {
			unsigned int buf_pos = 0;
			memcpy(ct, packet->buf, header_size - 2);
			*(unsigned short *)&ct[header_size - 2] = htons(length);
			tls_random(buf, TLS_AES_IV_LENGTH);
			buf_pos += TLS_AES_IV_LENGTH;
			/* copy payload */
			memcpy(buf + buf_pos, packet-> buf + header_size, packet->len - header_size);
			buf_pos += packet->len - header_size;

			tls_hmac_message(1, context,
					packet->buf,
					packet->len, NULL, 0,
					buf + buf_pos,
					mac_size);
			buf_pos += mac_size;

			memset(buf + buf_pos, padding - 1, padding);
			buf_pos += padding;

			crypto_encrypt(context, buf, ct + header_size, length);
			free(packet->buf);
			packet->buf = ct;
			packet->len = length + header_size;
			packet->size = packet->len;
		} else {
			/* invalidate packet */
			memset (packet->buf, 0, packet-> len);
		}
		free(buf);
	} else {
		/* invalidate packet */
		memset(packet->buf, 0, packet->len);
	}

	return 1;
}

static void put16(unsigned char *at, uint16_t val) {
	at[0] = val >> 8;
	at[1] = val & 0xff;
}

void tls_packet_update(struct TLSPacket *packet) {
	struct TLSContext *context;
	unsigned int header_size = 5;
	int footer_size = 0;
	uint16_t real_packet_len;
	uint64_t lsn;
	int mac_size = 0;
	unsigned int length = 0;
	int context_is_v13 = 0;
	struct tls_buffer ciphertext;

	if (!packet || packet->broken) {
		return;
	}

	if (packet->context && packet->context->tlsver == TLS_VERSION13
			&& packet->context->cipher_spec_set
			&& packet->context->crypto.created) {
		/* type */
		tls_packet_uint8(packet, packet->buf[0]);
		/* no padding
		 * tls_packet_uint8(packet, 0);
		 */
		footer_size = 1;
	}

	real_packet_len = packet->len - header_size;

	put16(packet->buf + 3, real_packet_len);

	if (!packet->context) {
		return;
	}

	context = packet->context;
	lsn = context->local_sequence_number;

	if (context->tlsver == TLS_VERSION13) {
		context_is_v13 = 1;
	}

	if (packet->buf[0] == TLS_CHANGE_CIPHER) {
		context->local_sequence_number++;
		return;
	}

	/* If this is a handshake message, update the handshake hash */
	if (packet->buf[0] == TLS_HANDSHAKE && packet->len > header_size) {
		unsigned char handshake_type = packet->buf[header_size];
		if (handshake_type != 0x00 && handshake_type != 0x03) {
			tls_update_hash(context, packet->buf + header_size,
					real_packet_len -
					footer_size);
		}
	}

	if (!context->cipher_spec_set || !context->crypto.created) {
		context->local_sequence_number++;
		return;
	}

	unsigned int pt_length = real_packet_len;

	if (context->crypto.created == 1) {
	} else if (context->crypto.created == 3) {
		mac_size = POLY1305_TAGLEN;
		length = real_packet_len + mac_size;
	} else {
		mac_size = TLS_GCM_TAG_LEN;
		length = real_packet_len + 8 + mac_size;
	}

	if (context->crypto.created == 1) {
		packet_encrypt(packet); 
		context->local_sequence_number++;
		return;
	}

	if (context->crypto.created < 1) {
		/* invalidate packet */
		memset(packet->buf, 0, packet->len);
		context->local_sequence_number++;
		return;
	}

	/* + 1 = type */
	int ct_size = length + header_size + 12 + TLS_MAX_TAG_LEN + 1;
	tls_buffer_init(&ciphertext, ct_size);

	if (ciphertext.error) {
		/* invalidate packet */
		memset(packet->buf, 0, packet->len);
		context->local_sequence_number++;
		return;
	}

	/* AEAD */
	/* sequence number (8 bytes) */
	/* content type (1 byte) */
	/* version (2 bytes) */
	/* length (2 bytes) */
	unsigned char aad[13];
	int aad_size = sizeof(aad);
	unsigned char *sequence = aad;
	if (context_is_v13) {
		aad[0] = TLS_APPLICATION_DATA;
		aad[1] = packet->buf[1];
		aad[2] = packet->buf[2];
		if (packet->context->crypto.created == 3) {
			put16(aad+3, real_packet_len + POLY1305_TAGLEN);
		} else {
			put16(aad+3, real_packet_len + TLS_GCM_TAG_LEN);
		}
		aad_size = 5;
		sequence = aad + 5;

		*((uint64_t *) (aad+5)) = htonll(lsn);

	} else {
		*((uint64_t *) aad) = htonll(lsn);
		aad[8] = packet->buf[0];
		aad[9] = packet->buf[1];
		aad[10] = packet->buf[2];
		put16(aad+11, packet->len - header_size);
	}

	ciphertext.len = header_size;

	if (context->crypto.created == 3) {
		int size;
		unsigned int counter = 1;
		unsigned char poly1305_key[POLY1305_KEYLEN];
		chacha_ivupdate(&context->crypto.ctx_local.chacha_local,
				context->crypto.ctx_local_mac.local_aead_iv,
				sequence, (uint8_t *) &counter);
		chacha20_poly1305_key(&context->crypto.ctx_local.chacha_local,
				poly1305_key);
		size =
			chacha20_poly1305_aead(&context->crypto.ctx_local.chacha_local,
					packet->buf + header_size, pt_length,
					aad, aad_size, poly1305_key,
					ciphertext.buffer + ciphertext.len);
		ciphertext.len += size;
	} else {
		unsigned char iv [TLS_13_AES_GCM_IV_LENGTH];
		if (context_is_v13) {
			memcpy(iv, context->crypto.ctx_local_mac.local_iv,
					TLS_13_AES_GCM_IV_LENGTH);
			int i;
			int offset = TLS_13_AES_GCM_IV_LENGTH - 8;
			for (i = 0; i < 8; i++)
				iv[offset + i] = context->crypto.ctx_local_mac.local_iv[offset + i] ^ sequence[i];
		} else {
			memcpy(iv, context->crypto.ctx_local_mac.local_aead_iv,
					TLS_AES_GCM_IV_LENGTH);
			tls_random(iv + TLS_AES_GCM_IV_LENGTH, 8);
			tls_buffer_append(&ciphertext,
					iv+TLS_AES_GCM_IV_LENGTH, 8);
		}

		gcm_state *localgcm;
		localgcm = &context->crypto.ctx_local.aes_gcm_local;

		gcm_reset(localgcm);
		gcm_add_iv(localgcm, iv, 12);
		gcm_add_aad(localgcm, aad, aad_size);
		gcm_process(localgcm, packet->buf + header_size, pt_length,
				ciphertext.buffer+ciphertext.len, GCM_ENCRYPT);
		ciphertext.len += pt_length;

		unsigned long taglen = TLS_GCM_TAG_LEN;
		gcm_done(localgcm, ciphertext.buffer+ciphertext.len, &taglen);
		ciphertext.len += taglen;
	}

	if (context_is_v13) {
		ciphertext.buffer[0] = TLS_APPLICATION_DATA;
		tls_buffer_write16(&ciphertext, TLS_V12, 1);
	} else {
		memcpy(ciphertext.buffer, packet->buf, header_size - 2);
	}

	tls_buffer_write16(&ciphertext, ciphertext.len - header_size, header_size - 2);

	free(packet->buf);
	packet->buf = ciphertext.buffer;
	packet->len = ciphertext.len;
	packet->size = ciphertext.size;

	context->local_sequence_number++;
}