#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <stdint.h>

#include <time.h>

#include <openssl/hmac.h>

#include "rr.h"
#include "utils.h"
#include "constants.h"
#include "packet.h"

/* ----- Creadores ----- */

DNSRR *dns2_rr_create_a (const char *name, int klass, int ttl, struct in_addr address) {
	DNSRR *rr;
	
	rr = (DNSRR *) malloc (sizeof (DNSRR));
	
	if (rr == NULL) {
		return NULL;
	}
	
	strncpy (rr->name, name, sizeof (rr->name));
	dns2_trim_name (rr->name);
	
	if (rr->name[0] == 0) {
		free (rr);
		return NULL;
	}
	
	rr->type = RR_TYPE_A;
	rr->klass = klass;
	rr->ttl = ttl;
	
	/* Copiar la IP */
	memcpy (&rr->a.address, &address, sizeof (rr->a.address));
	
	return rr;
}

DNSRR *dns2_rr_create_tsig (const char *name, const char *signature, const char *algorithm) {
	DNSRR *rr;
	
	if (strcmp (algorithm, RR_TSIG_HMAC_MD5) != 0 &&
	    strcmp (algorithm, RR_TSIG_HMAC_SHA1) != 0 &&
	    strcmp (algorithm, RR_TSIG_HMAC_SHA224) != 0 &&
	    strcmp (algorithm, RR_TSIG_HMAC_SHA256) != 0 &&
	    strcmp (algorithm, RR_TSIG_HMAC_SHA384) != 0 &&
	    strcmp (algorithm, RR_TSIG_HMAC_SHA512) != 0) {
		/* Algoritmo no soportado */
		return NULL;
	}
	    
	rr = (DNSRR *) malloc (sizeof (DNSRR));
	
	if (rr == NULL) {
		return NULL;
	}
	
	strncpy (rr->name, name, sizeof (rr->name));
	dns2_trim_name (rr->name);
	
	if (rr->name[0] == 0) {
		free (rr);
		return NULL;
	}
	
	rr->type = RR_TYPE_TSIG;
	rr->klass = RR_CLASS_ANY;
	rr->ttl = 0;
	
	strncpy (rr->tsig.algorithm, algorithm, sizeof (rr->tsig.algorithm));
	rr->tsig.key = strdup (signature);
	
	if (rr->tsig.key == NULL) {
		free (rr->name);
		free (rr);
		
		return NULL;
	}
	
	// El tiempo tiene que ser UTC
	// El tiempo son 6 bytes.
	rr->tsig.time_signed = time (NULL);
	//exit (1);
	//rr->tsig.time_signed = 1602910230u;
	//printf ("Time: %i\n", ((unsigned int) rr->tsig.time_signed));
	rr->tsig.fudge = 300;
	rr->tsig.mac_size = 0;
	rr->tsig.original_id = 0;
	
	rr->tsig.error = 0;
	rr->tsig.other_length = 0;
	
	return rr;
}

DNSRR *dns2_rr_create_aaaa (const char *name, int klass, int ttl, struct in6_addr address) {
	DNSRR *rr;
	
	rr = (DNSRR *) malloc (sizeof (DNSRR));
	
	if (rr == NULL) {
		return NULL;
	}
	
	strncpy (rr->name, name, sizeof (rr->name));
	dns2_trim_name (rr->name);
	
	if (rr->name[0] == 0) {
		free (rr);
		return NULL;
	}
	
	rr->type = RR_TYPE_AAAA;
	rr->klass = klass;
	rr->ttl = ttl;
	
	/* Copiar la IP */
	memcpy (&rr->aaaa.address, &address, sizeof (rr->aaaa.address));
	
	return rr;
}

DNSRR *dns2_rr_create_ns (const char *name, int klass, int ttl, const char *nsname) {
	DNSRR *rr;
	
	rr = (DNSRR *) malloc (sizeof (DNSRR));
	
	if (rr == NULL) {
		return NULL;
	}
	
	strncpy (rr->name, name, sizeof (rr->name));
	dns2_trim_name (rr->name);
	
	if (rr->name[0] == 0) {
		free (rr);
		return NULL;
	}
	
	rr->type = RR_TYPE_NS;
	rr->klass = klass;
	rr->ttl = ttl;
	
	/* Copiar el nombre */
	memcpy (&rr->ns.name, &nsname, strlen (nsname) + 1);
	dns2_trim_name (rr->ns.name);
	
	return rr;
}

/* ----- Empaquetadores ----- */
void dns_rr_pack_a_data (DNSRR *rr, DNSPacket *packet) {
	char *buf;
	uint16_t t16;
	
	buf = &packet->rdata[packet->rdlength];
	
	t16 = htons (4);
	memcpy (buf, &t16, 2);
	
	memcpy (&buf[2], &rr->a.address, 4);
	packet->rdlength += 4 + 2;
}

void dns_rr_pack_aaaa_data (DNSRR *rr, DNSPacket *packet) {
	char *buf;
	uint16_t t16;
	
	buf = &packet->rdata[packet->rdlength];
	
	t16 = htons (16);
	memcpy (buf, &t16, 2);
	
	memcpy (&buf[2], &rr->aaaa.address, 16);
	packet->rdlength += 16 + 2;
}

int dns_rr_sign_hmac (void *data, int data_len, const char *key, const char *algorithm, void *mac) {
	HMAC_CTX *ctx;
	const EVP_MD *md_algo;
	char *decoded_key;
	size_t key_length;
	int total_len;
	int g;
	
	printf ("About to sign, data is:\n");
	for (g = 0; g < data_len; g++) {
		printf ("%02x", (unsigned int) (((unsigned char *) data)[g]));
	}
	printf ("\nEnd data sig\n");
	
	if (strcmp (algorithm, RR_TSIG_HMAC_MD5) == 0) {
		md_algo = EVP_md5 ();
	} else if (strcmp (algorithm, RR_TSIG_HMAC_SHA1) == 0) {
		md_algo = EVP_sha1 ();
	} else if (strcmp (algorithm, RR_TSIG_HMAC_SHA224) == 0) {
		md_algo = EVP_sha224 ();
	} else if (strcmp (algorithm, RR_TSIG_HMAC_SHA256) == 0) {
		md_algo = EVP_sha256 ();
	} else if (strcmp (algorithm, RR_TSIG_HMAC_SHA384) == 0) {
		md_algo = EVP_sha384 ();
	} else if (strcmp (algorithm, RR_TSIG_HMAC_SHA512) == 0) {
		md_algo = EVP_sha512 ();
	} else {
		return 0;
	}
	
	ctx = HMAC_CTX_new ();
	
	if (ctx == NULL) {
		return 0;
	}
	
	if (base64_decode (key, (unsigned char **)&decoded_key, &key_length) != 0) {
		HMAC_CTX_free (ctx);
		
		return 0;
	}
	
	HMAC_Init_ex (ctx, decoded_key, key_length, md_algo, NULL);
	
	/* Enviar el mensaje a firmar */
	HMAC_Update (ctx, data, data_len);
	
	HMAC_Final (ctx, mac, &total_len);
	
	HMAC_CTX_free (ctx);
	
	free (decoded_key);
	
	return total_len;
}

void dns_rr_pack_tsig_data (DNSRR *rr, DNSPacket *packet) {
	uint16_t offset_rdlength, t16;
	uint32_t t32;
	
	/* Tomar la posición actual para el rdlength */
	offset_rdlength = packet->rdlength;
	packet->rdlength += 2;
	
	/* Agregar el algoritmo sin compresión */
	dns_packet_pack_name (packet, rr->tsig.algorithm);
	
	/* Empaquetar el tiempo, fudge y mac_size */
	t16 = 0;
	memcpy (&packet->rdata[packet->rdlength], &t16, 2);
	packet->rdlength += 2;
	
	t32 = htonl ((intmax_t) rr->tsig.time_signed);
	memcpy (&packet->rdata[packet->rdlength], &t32, 4);
	packet->rdlength += 4;
	
	/* El fudge */
	t16 = htons (rr->tsig.fudge);
	memcpy (&packet->rdata[packet->rdlength], &t16, 2);
	packet->rdlength += 2;
	
	t16 = htons (rr->tsig.mac_size);
	memcpy (&packet->rdata[packet->rdlength], &t16, 2);
	packet->rdlength += 2;
	
	/* Agregar el MAC */
	memcpy (&packet->rdata[packet->rdlength], rr->tsig.mac, rr->tsig.mac_size);
	packet->rdlength += rr->tsig.mac_size;
	
	/* Copiar el packet id */
	t16 = htons (packet->header.id);
	memcpy (&packet->rdata[packet->rdlength], &t16, 2);
	packet->rdlength += 2;
	
	/* Error y other length */
	t16 = htons (rr->tsig.error);
	memcpy (&packet->rdata[packet->rdlength], &t16, 2);
	packet->rdlength += 2;
	
	t16 = htons (rr->tsig.other_length);
	memcpy (&packet->rdata[packet->rdlength], &t16, 2);
	packet->rdlength += 2;
	
	/* Agregar el other_data si other_length > 0 */
	if (rr->tsig.other_length > 0) {
		memcpy (&packet->rdata[packet->rdlength], rr->tsig.other_data, rr->tsig.other_length);
		packet->rdlength += rr->tsig.other_length;
	}
	
	/* Regresar a la posición del rdlength y escribir el valor correcto */
	t16 = htons (packet->rdlength - offset_rdlength - 2);
	memcpy (&packet->rdata[offset_rdlength], &t16, 2);
}

void dns_rr_pack_generate_signature (DNSRR *signature, DNSPacket *packet) {
	void *data;
	int data_count;
	uint16_t t16;
	uint32_t t32;
	
	/* Reempaquetar todo el paquete para generar los bytes de la firma */
	dns_packet_reset (packet);
	dns_packet_pack_data (packet);
	
	/* Agregar el nombre del tsig sin compresión */
	dns_packet_pack_name (packet, signature->name);
	
	t16 = htons (signature->klass);
	memcpy (&packet->rdata[packet->rdlength], &t16, 2);
	packet->rdlength += 2;
	
	t32 = htonl (signature->ttl);
	memcpy (&packet->rdata[packet->rdlength], &t32, 4);
	packet->rdlength += 4;
	
	/* Agregar el algoritmo sin compresión */
	dns_packet_pack_name (packet, signature->tsig.algorithm);
	
	/* Agregar la hora */
	t16 = 0;
	memcpy (&packet->rdata[packet->rdlength], &t16, 2);
	packet->rdlength += 2;
	
	t32 = htonl ((intmax_t) signature->tsig.time_signed);
	memcpy (&packet->rdata[packet->rdlength], &t32, 4);
	packet->rdlength += 4;
	
	/* El fudge */
	t16 = htons (signature->tsig.fudge);
	memcpy (&packet->rdata[packet->rdlength], &t16, 2);
	packet->rdlength += 2;
	
	/* Error y other length */
	t16 = htons (signature->tsig.error);
	memcpy (&packet->rdata[packet->rdlength], &t16, 2);
	packet->rdlength += 2;
	t16 = htons (signature->tsig.other_length);
	memcpy (&packet->rdata[packet->rdlength], &t16, 2);
	packet->rdlength += 2;
	
	/* Agregar el other_data si other_length > 0 */
	if (signature->tsig.other_length > 0) {
		memcpy (&packet->rdata[packet->rdlength], signature->tsig.other_data, signature->tsig.other_length);
		packet->rdlength += signature->tsig.other_length;
	}
	
	/* Recuperar la data */
	data_count = dns_packet_get_data (packet, (unsigned char **)&data);
	 
	/* Realmente firmar la data */
	signature->tsig.mac_size = dns_rr_sign_hmac (data, data_count, signature->tsig.key, signature->tsig.algorithm, signature->tsig.mac);
	
	/* Pedir reempaquetar la data, otra vez */
	dns_packet_reset (packet);
}

void dns_rr_pack_data (DNSRR *rr, DNSPacket *packet) {
	uint16_t t16;
	uint32_t t32;
	char *buf;
	
	dns_packet_compress_name (packet, rr->name);
	
	buf = &packet->rdata[packet->rdlength];
	
	/*if (rr->type == OPT) {
	
	} else { */
		t16 = htons (rr->type);
		memcpy (&buf[0], &t16, 2);
		
		t16 = htons (rr->klass);
		memcpy (&buf[2], &t16, 2);
		
		t32 = htonl (rr->ttl);
		memcpy (&buf[4], &t32, 4);
	/* } */
	
	packet->rdlength += 8;
	
	/* Data por tipo */
	switch (rr->type) {
		case RR_TYPE_A:
			dns_rr_pack_a_data (rr, packet);
			break;
		case RR_TYPE_AAAA:
			dns_rr_pack_aaaa_data (rr, packet);
			break;
		case RR_TYPE_TSIG:
			dns_rr_pack_tsig_data (rr, packet);
			break;
	}
}

/* ----- Desempaquetadores ----- */
int dns_rr_unpack_a_data (DNSRR *rr, unsigned char *buf, int rd_length) {
	if (rd_length < 4) {
		return -1;
	}
	
	memcpy (&rr->a.address, buf, 4);
	
	return 0;
}

int dns_rr_unpack_aaaa_data (DNSRR *rr, unsigned char *buf, int rd_length) {
	if (rd_length < 16) {
		return -1;
	}
	
	memcpy (&rr->aaaa.address, buf, 16);
	
	return 0;
}

int dns_rr_unpack_ns_data (DNSRR *rr, unsigned char *buf, int rd_length, DNSPacket *packet) {
	/* No hacemos validaciones sobre la longitud porque los nombres son variables */
	int save;
	save = buf - packet->rdata;
	
	if (dns_packet_expand_name (rr->ns.name, packet, &save) < 0) {
		return -1;
	}
	
	return 0;
}

int dns_rr_unpack_cname_data (DNSRR *rr, unsigned char *buf, int rd_length, DNSPacket *packet) {
	/* No hacemos validaciones sobre la longitud porque los nombres son variables */
	int save;
	save = buf - packet->rdata;
	
	if (dns_packet_expand_name (rr->cname.name, packet, &save) < 0) {
		return -1;
	}
	
	return 0;
}

int dns_rr_unpack_soa_data (DNSRR *rr, unsigned char *buf, int rd_length, DNSPacket *packet) {
	int save, start;
	int tot;
	uint32_t u32;
	
	if (rd_length < 20) { /* Al menos 5 enteros de 32 bits cada uno */
		return -1;
	}
	
	start = save = buf - packet->rdata;
	if (dns_packet_expand_name (rr->soa.mname, packet, &save) < 0) {
		return -1;
	}
	
	if (dns_packet_expand_name (rr->soa.rname, packet, &save) < 0) {
		return -1;
	}
	
	tot = save - start;
	if (rd_length - tot < 20) {
		return -1;
	}
	
	/* Leer el serial */
	memcpy (&u32, &buf[tot], 4);
	rr->soa.serial = ntohl (u32);
	
	/* Leer el refresh */
	memcpy (&u32, &buf[tot + 4], 4);
	rr->soa.refresh = ntohl (u32);
	
	/* Leer el retry */
	memcpy (&u32, &buf[tot + 8], 4);
	rr->soa.retry = ntohl (u32);
	
	/* El expire */
	memcpy (&u32, &buf[tot + 12], 4);
	rr->soa.expire = ntohl (u32);
	
	/* El minimum */
	memcpy (&u32, &buf[tot + 16], 4);
	rr->soa.minimum = ntohl (u32);
	
	return 0;
}

int dns_rr_unpack_ptr_data (DNSRR *rr, unsigned char *buf, int rd_length, DNSPacket *packet) {
	int save;
	save = buf - packet->rdata;
	
	if (dns_packet_expand_name (rr->ptr.name, packet, &save) < 0) {
		return -1;
	}
	
	return 0;
}

int dns_rr_unpack_mx_data (DNSRR *rr, unsigned char *buf, int rd_length, DNSPacket *packet) {
	int save;
	uint16_t u16;
	
	if (rd_length < 2) {
		return -1;
	}
	
	memcpy (&u16, buf, 2);
	rr->mx.preference = ntohs (u16);
	
	save = &buf[2] - packet->rdata;
	
	if (dns_packet_expand_name (rr->mx.name, packet, &save) < 0) {
		return -1;
	}
	
	return 0;
}

int dns_rr_unpack_txt_data (DNSRR *rr, unsigned char *buf, int rd_length) {
	/* En la teoría, un TXT puede contener muchas cadenas,
	 * En la práctica, no creo que ocurra.
	 * Tomar solo la primer cadena */
	uint8_t txt_len;
	
	txt_len = buf[0];
	
	if (txt_len > rd_length - 1) {
		return -1;
	}
	memcpy (rr->txt.txt, &buf[1], txt_len);
	
	rr->txt.txt[txt_len] = 0;
	
	return 0;
}

int dns_rr_unpack_tsig_data (DNSRR *rr, unsigned char *buf, int rd_length) {
	uint32_t t32;
	uint16_t t16;
	uint8_t t8;
	
	int offset;
	
	offset = 0;
	t8 = dns_packet_expand_name_simple (rr->tsig.algorithm, buf, rd_length);
	
	if (t8 < 0) {
		return -1;
	}
	
	offset += t8;
	
	/* Revisar el time, fudge y mac_size */
	if (offset + 10 > rd_length) {
		return -1;
	}
	
	/* ¿Ignorar los primeros dos bytes del tiempo? */
	memcpy (&t16, &buf[offset], 2);
	
	memcpy (&t32, &buf[offset + 2], 4);
	rr->tsig.time_signed = ntohl (t32);
	
	memcpy (&t16, &buf[offset + 6], 2);
	rr->tsig.fudge = ntohs (t16);
	
	memcpy (&t16, &buf[offset + 8], 2);
	rr->tsig.mac_size = ntohs (t16);
	
	offset += 10;
	
	/* Copiar el MAC */
	if (offset + rr->tsig.mac_size > rd_length) {
		return -1;
	}
	
	memcpy (&rr->tsig.mac, &buf[offset], rr->tsig.mac_size);
	offset += rr->tsig.mac_size;
	
	/* Revisar el packet id, error y other_len */
	if (offset + 6 > rd_length) {
		return -1;
	}
	
	/* Copiar el packet id original */
	memcpy (&t16, &buf[offset], 2);
	rr->tsig.original_id = ntohs (t16);
	
	memcpy (&t16, &buf[offset + 2], 2);
	rr->tsig.error = ntohs (t16);
	
	memcpy (&t16, &buf[offset + 4], 2);
	rr->tsig.other_length = ntohs (t16);
	
	offset += 6;
	if (rr->tsig.other_length > 0) {
		if (offset + rr->tsig.other_length < rd_length) {
			return -1;
		}
		
		memcpy (&rr->tsig.other_data, &buf[offset], rr->tsig.other_length);
	}
	
	return 0;
}

DNSRR * dns_rr_unpack_data (DNSPacket *packet) {
	DNSRR *rr;
	uint16_t t16;
	uint32_t t32;
	unsigned char *buf;
	int rd_length;
	int res;
	
	rr = (DNSRR *) malloc (sizeof (DNSRR));
	
	if (rr == NULL) {
		return NULL;
	}
	
	memset (rr, 0, sizeof (DNSRR));
	
	if (dns_packet_expand_name (rr->name, packet, &packet->rdread) < 0) {
		dns2_rr_free (rr);
		
		return NULL;
	}
	
	/* Asegurar al menos 10 bytes */
	if (packet->rdlength < packet->rdread + 10) {
		dns2_rr_free (rr);
		
		return NULL;
	}
	
	buf = &packet->rdata[packet->rdread];
	packet->rdread += 10;
	
	memcpy (&t16, &buf[0], sizeof (t16));
	rr->type = ntohs (t16);
	
	memcpy (&t16, &buf[2], sizeof (t16));
	rr->klass = ntohs (t16);
	
	memcpy (&t32, &buf[4], sizeof (t32));
	rr->ttl = ntohl (t32);
	
	memcpy (&t16, &buf[8], sizeof (t16));
	rd_length = ntohs (t16);
	
	if (packet->rdlength < packet->rdread + rd_length) {
		dns2_rr_free (rr);
		
		return NULL;
	}
	
	buf = &packet->rdata[packet->rdread];
	packet->rdread += rd_length;
	
	res = 0;
	switch (rr->type) {
		case RR_TYPE_A:
			res = dns_rr_unpack_a_data (rr, buf, rd_length);
			break;
		case RR_TYPE_AAAA:
			res = dns_rr_unpack_aaaa_data (rr, buf, rd_length);
			break;
		case RR_TYPE_NS:
			res = dns_rr_unpack_ns_data (rr, buf, rd_length, packet);
			break;
		case RR_TYPE_CNAME:
			res = dns_rr_unpack_cname_data (rr, buf, rd_length, packet);
			break;
		case RR_TYPE_SOA:
			res = dns_rr_unpack_soa_data (rr, buf, rd_length, packet);
			break;
		case RR_TYPE_PTR:
			res = dns_rr_unpack_ptr_data (rr, buf, rd_length, packet);
			break;
		case RR_TYPE_MX:
			res = dns_rr_unpack_mx_data (rr, buf, rd_length, packet);
			break;
		case RR_TYPE_TXT:
			res = dns_rr_unpack_txt_data (rr, buf, rd_length);
			break;
		case RR_TYPE_TSIG:
			res = dns_rr_unpack_tsig_data (rr, buf, rd_length);
			break;
	}
	
	if (res < 0) {
		dns2_rr_free (rr);
		
		return NULL;
	}
	
	return rr;
}

DNSRR *dns2_rr_dup (const DNSRR *rr) {
	DNSRR *new;
	
	new = (DNSRR *) malloc (sizeof (DNSRR));
	if (new == NULL) {
		return NULL;
	}
	
	memcpy (new, rr, sizeof (DNSRR));
	
	if (rr->type == RR_TYPE_TSIG) {
		new->tsig.key = strdup (rr->tsig.key);
	}
	
	return new;
}

void dns2_rr_free (DNSRR *rr) {
	if (rr->type == RR_TYPE_TSIG) {
		/* La llave se duplicó, liberarla */
		free (rr->tsig.key);
	}
	
	free (rr);
}