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

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

void dns_rr_pack_data (DNSRR *rr, DNSPacket *packet);
DNSRR * dns_rr_unpack_data (DNSPacket *packet);

DNSPacket *dns_packet_create_request (const char *name, int type, int klass) {
	DNSPacket *req;
	DNSQuestion *q;
	
	req = (DNSPacket *) malloc (sizeof (DNSPacket));
	
	if (req == NULL) return NULL;
	
	dns_header_init (&req->header);
	
	q = dns_question_create (name, type, klass);
	
	if (q == NULL) {
		free (req);
		
		return NULL;
	}
	
	req->questions = g_list_append (NULL, q);
	
	req->authority = NULL;
	req->answers = NULL;
	req->additional = NULL;
	req->rdlength = 0;
	
	req->labels = NULL;
	
	return req;
}

DNSPacket *dns_packet_create_response (int type, unsigned char *data, int size_data) {
	DNSPacket *res;
	DNSRR *rr;
	
	int g;
	
	res = (DNSPacket *) malloc (sizeof (DNSPacket));
	
	if (res == NULL) return NULL;
	
	memset (res, 0, sizeof (DNSPacket));
	
	res->answer_socket_type = type;
	
	memcpy (res->rdata, data, size_data);
	res->rdlength = size_data;
	res->rdread = 0;
	
	if (dns_header_init_from_packet (&res->header, res) < 0) {
		dns_packet_free_full (res);
		
		return NULL;
	}
	
	/* TODO: Revisar el bit de truncamiento */
	
	/* Parsear las preguntas */
	for (g = 0; g < res->header.qdcount; g++) {
		if (dns_question_parse_question_from_packet (res) < 0) {
			dns_packet_free_full (res);
			
			return NULL;
		}
	}
	
	/* Parsear los answers */
	for (g = 0; g < res->header.ancount; g++) {
		rr = dns_rr_unpack_data (res);
		
		if (rr == NULL) {
			dns_packet_free_full (res);
			
			return NULL;
		}
		
		res->answers = g_list_append (res->answers, rr);
	}
	
	/* Parsear los authority */
	for (g = 0; g < res->header.nscount; g++) {
		rr = dns_rr_unpack_data (res);
		
		if (rr == NULL) {
			dns_packet_free_full (res);
			
			return NULL;
		}
		
		res->authority = g_list_append (res->authority, rr);
	}
	
	/* Parsear los aditional */
	for (g = 0; g < res->header.arcount; g++) {
		rr = dns_rr_unpack_data (res);
		
		if (rr == NULL) {
			dns_packet_free_full (res);
			
			return NULL;
		}
		
		res->additional = g_list_append (res->additional, rr);
	}
	
	dns_packet_recount_items (res);
	
	return res;
}

int dns_packet_get_data (DNSPacket *packet, unsigned char **data) {
	if (data != NULL) {
		*data = packet->rdata;
	}
	
	return packet->rdlength;
}

void dns_packet_pack_data (DNSPacket *packet) {
	GList *g;
	DNSQuestion *qq;
	DNSRR *rr;
	
	/* Reiniciar desde 0 */
	packet->rdlength = 0;
	
	dns_header_pack_data (&packet->header, packet);
	
	g = packet->questions;
	while (g != NULL) {
		qq = (DNSQuestion *) g->data;
		
		dns_question_pack_data (qq, packet);
		
		g = g->next;
	}
	
	g = packet->answers;
	while (g != NULL) {
		rr = (DNSRR *) g->data;
		
		dns_rr_pack_data (rr, packet);
		
		g = g->next;
	}
	
	g = packet->authority;
	while (g != NULL) {
		rr = (DNSRR *) g->data;
		
		dns_rr_pack_data (rr, packet);
		
		g = g->next;
	}
	
	g = packet->additional;
	while (g != NULL) {
		rr = (DNSRR *) g->data;
		
		dns_rr_pack_data (rr, packet);
		
		g = g->next;
	}
}

static int dns_packet_search_label (gconstpointer left, gconstpointer right) {
	PacketLabel *a = (PacketLabel *) left;
	PacketLabel *b = (PacketLabel *) right;
	
	return strcmp (a->name, b->name);
}

static void dns_packet_append_word (DNSPacket *packet, const char *name) {
	const char *ending;
	PacketLabel *new_label;
	uint8_t t8;
	
	ending = strchr (name, '.');
	if (ending == NULL) {
		ending = name + strlen (name);
	}
	
	new_label = (PacketLabel *) malloc (sizeof (PacketLabel));
		
	if (new_label == NULL) {
		/* Oops, error */
		return;
	}
	
	strncpy (new_label->name, name, sizeof (new_label->name));
	new_label->offset = packet->rdlength;
	
	packet->labels = g_list_append (packet->labels, new_label);
	
	/* Copiar solo la palabra */
	t8 = (ending - name);
	packet->rdata[packet->rdlength] = t8;
	memcpy (&packet->rdata[packet->rdlength + 1], name, t8);
	
	packet->rdlength += t8 + 1;
}

void dns_packet_pack_name (DNSPacket *packet, const char *name) {
	char * dot_pos;
	
	while (dot_pos = strchr (name, '.'), dot_pos != NULL) {
		/* Buscar esta palabra en nuestro diccionario de palabras */
		name = dot_pos + 1;
		
		/* Agregar esta palabra al diccionario de etiquetas */
		dns_packet_append_word (packet, name);
	}
	
	/* La palabra restante, copiarla */
	if (name[0] != 0) {
		dns_packet_append_word (packet, name);
	}
	
	packet->rdata[packet->rdlength] = 0;
	packet->rdlength += 1;
}

void dns_packet_compress_name (DNSPacket *packet, char *name) {
	PacketLabel temp, *found;
	char * dot_pos;
	GList *search;
	uint16_t t16;
	
	while (dot_pos = strchr (name, '.'), dot_pos != NULL) {
		/* Buscar esta palabra en nuestro diccionario de palabras */
		
		strncpy (temp.name, name, sizeof (temp.name));
		
		search = g_list_find_custom (packet->labels, &temp, dns_packet_search_label);
		
		if (search != NULL) {
			found = (PacketLabel *) search->data;
			
			t16 = htons (0xc000 | found->offset);
			
			memcpy (&packet->rdata[packet->rdlength], &t16, 2);
			packet->rdlength += 2;
			
			return;
		}
		
		/* Agregar esta palabra al direccionario de etiquetas */
		dns_packet_append_word (packet, name);
		name = dot_pos + 1;
	}
	
	/* La palabra restante, copiarla */
	if (name[0] != 0) {
		dns_packet_append_word (packet, name);
	}
	
	packet->rdata[packet->rdlength] = 0;
	packet->rdlength += 1;
}

DNSPacket *dns_packet_clone_simple (DNSPacket *packet) {
	DNSPacket *clone;
	
	clone = (DNSPacket *) malloc (sizeof (DNSPacket));
	
	if (clone == NULL) return NULL;
	
	memcpy (&clone->header, &packet->header, sizeof (DNSHeader));
	
	clone->questions = g_list_copy (packet->questions);
	clone->answers = g_list_copy (packet->answers);
	clone->authority = g_list_copy (packet->authority);
	clone->additional = g_list_copy (packet->additional);
	
	clone->labels = NULL;
	
	clone->rdlength = 0;
	
	/* NOTA: No se copian las labels */
	
	return clone;
}

void dns_packet_free_simple (DNSPacket *packet) {
	/* Liberar las estructuras, excepto los datos que no duplicamos */
	g_list_free_full (packet->labels, (GDestroyNotify) free);
	
	g_list_free (packet->questions);
	g_list_free (packet->answers);
	g_list_free (packet->authority);
	g_list_free (packet->additional);
	
	free (packet);
}

void dns_packet_free_full (DNSPacket *packet) {
	g_list_free_full (packet->labels, (GDestroyNotify) free);
	
	g_list_free_full (packet->questions, (GDestroyNotify) dns_question_free);
	g_list_free_full (packet->answers, (GDestroyNotify) dns2_rr_free);
	g_list_free_full (packet->authority, (GDestroyNotify) dns2_rr_free);
	g_list_free_full (packet->additional, (GDestroyNotify) dns2_rr_free);
	
	free (packet);
}

void dns_packet_reset (DNSPacket *packet) {
	g_list_free_full (packet->labels, (GDestroyNotify) free);
	packet->labels = NULL;
	
	packet->rdlength = 0;
}

int dns_packet_expand_name (unsigned char *name, DNSPacket *packet, int *offset) {
	int name_len = 0;
	uint8_t t8;
	uint16_t t16;
	int local_offset;
	int ptr, ptr_inicio;
	int g;
	
	while (1) {
		if (name_len >= 256) {
			return -1;
		}
		
		if (packet->rdlength < (*offset) + 1) {
			return -1;
		}
		
		t8 = packet->rdata[*offset];
		
		if (t8 == 0) {
			name[name_len] = 0;
			
			(*offset)++;
			
			return 0;
		} else if ((t8 & 0xc0) == 0xc0) {
			/* Tenemos un apuntador */
			if (packet->rdlength < (*offset) + 2) {
				return -1;
			}
			
			memcpy (&t16, &packet->rdata[*offset], sizeof (t16));
			t16 = ntohs (t16);
			
			ptr = t16 & 0x3fff;
			ptr_inicio = ptr;
			if (dns_packet_expand_name (&name[name_len], packet, &ptr) < 0) {
				return -1;
			}
			
			/* Sumar a la longitud de nombres lo recibido del apuntador */
			name_len = name_len + (ptr - ptr_inicio);
			(*offset) += 2;
			
			break;
		} else {
			(*offset)++;
			
			if (packet->rdlength < (*offset) + t8) {
				return -1;
			}
			
			/* Una palabra */
			for (g = *offset; g < t8 + *offset; g++) {
				name[name_len] = packet->rdata[g];
				name_len++;
			}
			
			name[name_len] = '.';
			name_len++;
			
			(*offset) += t8;
		}
	}
	
	/* Limpiar el último punto */
	if (name[name_len - 1] == '.') {
		name[name_len - 1] = 0;
		name_len--;
	}
	
	return 0;
}

int dns_packet_expand_name_simple (unsigned char *name, unsigned char *buffer, int buffer_size) {
	int name_len = 0;
	uint8_t t8;
	int g;
	int offset = 0;
	
	while (1) {
		if (name_len >= 256) {
			return -1;
		}
		
		if (buffer_size < offset + 1) {
			return -1;
		}
		
		t8 = buffer[offset];
		offset++;
		
		if (t8 == 0) {
			/* Fin del nombre */
			name[name_len] = 0;
			
			return offset;
		} else {
			if (buffer_size < offset + t8) {
				return -1;
			}
			
			/* Una palabra */
			for (g = offset; g < t8 + offset; g++) {
				name[name_len] = buffer[g];
				name_len++;
			}
			
			name[name_len] = '.';
			name_len++;
			
			offset += t8;
		}
	}
	
	/* Limpiar el último punto */
	if (name[name_len - 1] == '.') {
		name[name_len - 1] = 0;
		name_len--;
	}
	
	return offset;
}

void dns_header_pack_data (DNSHeader *header, DNSPacket *packet) {
	unsigned char *buf;
	uint16_t t16;
	uint8_t t8;
	
	buf = &packet->rdata[packet->rdlength];
	
	/* Copiar el ID */
	t16 = htons (header->id);
	memcpy (&buf[0], &t16, sizeof (t16));
	
	t8 = (header->qr << 7) | (header->opcode << 3) | (header->aa << 2) | (header->tc << 1) | (header->rd);
	buf[2] = t8;
	
	t8 = (header->ra << 7) | (header->ad << 5) | (header->cd << 4) | header->rcode;
	buf[3] = t8;
	
	t16 = htons (header->qdcount);
	memcpy (&buf[4], &t16, sizeof (t16));
	
	t16 = htons (header->ancount);
	memcpy (&buf[6], &t16, sizeof (t16));
	
	t16 = htons (header->nscount);
	memcpy (&buf[8], &t16, sizeof (t16));
	
	t16 = htons (header->arcount);
	memcpy (&buf[10], &t16, sizeof (t16));
	
	packet->rdlength += DNS_HEADER_SIZE;
}

int dns_header_init_from_packet (DNSHeader *header, DNSPacket *packet) {
	uint16_t t16;
	uint8_t t8;
	unsigned char *buf;
	
	buf = &packet->rdata[packet->rdread];
	
	if (packet->rdlength < packet->rdread + DNS_HEADER_SIZE) {
		return -1;
	}
	
	memcpy (&t16, &buf[0], sizeof (t16));
	header->id = ntohs (t16);
	
	t8 = buf[2];
	
	header->qr = (t8 >> 7) & 0x1;
	header->opcode = (t8 >> 3) & 0xf;
	header->aa = (t8 >> 2) & 0x1;
	header->tc = (t8 >> 1) & 0x1;
	header->rd = t8 & 0x1;
	
	t8 = buf[3];
	header->ra = (t8 >> 7) & 0x1;
	header->z = (t8 >> 6) & 0x1;
	header->ad = (t8 >> 5) & 0x1;
	header->cd = (t8 >> 4) & 0x1;
	header->rcode = t8 & 0xf;
	
	memcpy (&t16, &buf[4], sizeof (t16));
	header->qdcount = ntohs (t16);
	
	memcpy (&t16, &buf[6], sizeof (t16));
	header->ancount = ntohs (t16);
	
	memcpy (&t16, &buf[8], sizeof (t16));
	header->nscount = ntohs (t16);
	
	memcpy (&t16, &buf[10], sizeof (t16));
	header->arcount = ntohs (t16);
	
	packet->rdread += DNS_HEADER_SIZE;
	
	return 0;
}

void dns_packet_recount_items (DNSPacket *packet) {
	if (packet == NULL) return;
	
	packet->header.qdcount = g_list_length (packet->questions);
	packet->header.ancount = g_list_length (packet->answers);
	packet->header.nscount = g_list_length (packet->authority);
	packet->header.arcount = g_list_length (packet->additional);
}

DNSRR **dns_packet_get_answers (DNSPacket *packet) {
	DNSRR **array;
	GList *g;
	int n;
	
	dns_packet_recount_items (packet);
	
	if (packet->header.ancount == 0) {
		return NULL;
	}
	
	array = (DNSRR **) malloc ((packet->header.ancount + 1) * sizeof (DNSRR *));
	
	for (n = 0, g = packet->answers; g != NULL; g = g->next, n++) {
		array[n] = (DNSRR *) g->data;
	}
	
	array[n] = NULL;
	
	return array;
}

DNSRR **dns_packet_get_authority (DNSPacket *packet) {
	DNSRR **array;
	GList *g;
	int n;
	
	dns_packet_recount_items (packet);
	
	if (packet->header.nscount == 0) {
		return NULL;
	}
	
	array = (DNSRR **) malloc ((packet->header.nscount + 1) * sizeof (DNSRR *));
	
	for (n = 0, g = packet->authority; g != NULL; g = g->next, n++) {
		array[n] = (DNSRR *) g->data;
	}
	
	array[n] = NULL;
	
	return array;
}

DNSRR **dns_packet_get_additional (DNSPacket *packet) {
	DNSRR **array;
	GList *g;
	int n;
	
	dns_packet_recount_items (packet);
	
	if (packet->header.arcount == 0) {
		return NULL;
	}
	
	array = (DNSRR **) malloc ((packet->header.arcount + 1) * sizeof (DNSRR *));
	
	for (n = 0, g = packet->additional; g != NULL; g = g->next, n++) {
		array[n] = (DNSRR *) g->data;
	}
	
	array[n] = NULL;
	
	return array;
}