/*
 * 2008+ Copyright (c) Evgeniy Polyakov <johnpol@2ka.mipt.ru>
 * All rights reserved.
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 */

#include <sys/types.h>
#include <sys/socket.h>

#include <netdb.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <dirent.h>
#include <pthread.h>

#include <netinet/in.h>
#include <arpa/inet.h>

#include "query.h"

static char query_spaces[] = "        ";

void query_fill_header(struct query_header *h, __u16 id)
{
	h->id = id;

	h->flags = QUERY_FLAGS_RD;
	query_set_flags_opcode(&h->flags, QUERY_FLAGS_OPCODE_STANDARD);
	query_set_flags_rcode(&h->flags, QUERY_FLAGS_RCODE_NOERROR);

	h->question_num = 0;
	h->answer_num = 0;
	h->auth_num = 0;
	h->addon_num = 0;
}

int query_fill_name(char *query, char *name)
{
	char *start = name;
	char *label;
	int len;

	while ((label = strchr(query, '.')) != NULL) {
		*label++ = '\0';
		len = strlen(query);
		*name++ = len;
		strncpy(name, query, len);
		name += len;
		query = label;

		if (*query == '\0')
			break;
	}

	len = strlen(query);
	*name++ = len;
	strncpy(name, query, len);
	name += len;
	*name++ = '\0';

	return name - start;
}

void *query_fill_question(void *q, char *query, __u16 type, __u16 class)
{
	__u16 *ptr;
	int len;

	len = query_fill_name(query, q);

	ptr = (__u16 *)(q + len);

	*ptr++ = htons(type);
	*ptr++ = htons(class);

	return ptr;
}

void query_header_convert(struct query_header *h)
{
	h->id = ntohs(h->id);
	h->flags = ntohs(h->flags);
	h->question_num = ntohs(h->question_num);
	h->answer_num = ntohs(h->answer_num);
	h->auth_num = ntohs(h->auth_num);
	h->addon_num = ntohs(h->addon_num);
}

void query_parse_header(struct query_header *h)
{
	__u16 opcode, rcode;

	query_header_convert(h);

	opcode = (h->flags >> QUERY_FLAGS_OPCODE_SHIFT) & QUERY_FLAGS_OPCODE_MASK;
	rcode = (h->flags >> QUERY_FLAGS_RCODE_SHIFT) & QUERY_FLAGS_RCODE_MASK;

	uloga("id: %04x: flags: resp: %d, opcode: %d, auth: %d, trunc: %d, RD: %d, RA: %d, rcode: %d.\n",
			h->id,
			!!(h->flags & QUERY_FLAGS_RESPONSE),
			opcode,
			!!(h->flags & QUERY_FLAGS_AA),
			!!(h->flags & QUERY_FLAGS_TC),
			!!(h->flags & QUERY_FLAGS_RD),
			!!(h->flags & QUERY_FLAGS_RA),
			rcode);
	uloga("%s: question: %d, answer: %d, auth: %d, addon: %d.\n",
			query_spaces, h->question_num, h->answer_num, h->auth_num, h->addon_num);
}

int query_parse_name(void *message, char *nptr, char *dst, int *off)
{
	unsigned char len;
	int err, name_len;

	len = 0;
	name_len = 0;
	*off = 0;

	while ((len = *nptr)) {
		if (len & 0xc0) {
			__u16 o = ((__u16 *)nptr)[0];
			__u16 offset = ntohs(o) & ~0xc000;

			err = query_parse_name(message, message + offset, dst, off);
			dst += err;
			name_len += err;
			nptr += 2;
			*off = 2;
			return name_len;
		} else {
			nptr++;
			strncpy(dst, nptr, len);
			dst += len;
			*dst++ = '.';
			nptr += len;
			name_len += len + 1;
		}
	}
	*dst = '\0';
	name_len++;

	*off = name_len;

	return name_len;
}

static int query_parse_rdata_ns(void *message, void *rdata,
		__u16 rlen __attribute__ ((unused)), char *dst,
		int dsize __attribute__ ((unused)))
{
	int offset;
	return query_parse_name(message, rdata, dst, &offset);
}

static int query_parse_rdata_a(void *message __attribute__ ((unused)),
		void *rdata, __u16 rlen, char *dst, int dsize)
{
	if (rlen != 4)
		return -EINVAL;

	return snprintf(dst, dsize, "%s", inet_ntoa(*((struct in_addr *)rdata)));
}

typedef int (* query_parse_rdata_type)(void *message, void *rdata, __u16 rlen, char *dst, int dsize);

static query_parse_rdata_type query_parse_rdata[256] = {
	[QUERY_TYPE_A] = query_parse_rdata_a,
	[QUERY_TYPE_NS] = query_parse_rdata_ns,
};

struct rr *query_parse_rr(void *message, void *data, unsigned int *off)
{
	struct rr *rr;
	__u16 rlen;
	int name_len, offset, header_inner_size = 10;
	char name[QUERY_RR_NAME_MAX_SIZE];

	name_len = query_parse_name(message, data, name, &offset);
	data += offset;

	rlen = ntohs(((__u16 *)data)[4]);

	rr = malloc(sizeof(struct rr) + rlen + 1);
	if (!rr)
		return NULL;

	memset(rr, 0, sizeof(rr) + rlen + 1);

	rr->namelen = snprintf(rr->name, sizeof(rr->name), "%s", name);
	rr->type = ntohs(((__u16 *)data)[0]);
	rr->class = ntohs(((__u16 *)data)[1]);
	rr->ttl = ntohl(((__u32 *)data)[1]);
	rr->rdlen = rlen;
	memcpy(rr->rdata, data + header_inner_size, rr->rdlen);

	uloga("%s: name: '%s', type: %d, class: %d, ttl: %u, rdlen: %d",
			query_spaces, rr->name, rr->type, rr->class, rr->ttl, rr->rdlen);

	if (query_parse_rdata[rr->type]) {
		char rdata[QUERY_RR_NAME_MAX_SIZE];

		query_parse_rdata[rr->type](message, rr->rdata, rr->rdlen, rdata, sizeof(rdata));
		uloga(", rdata: %s", rdata);
	}

	uloga("\n");

	*off = offset + rlen + header_inner_size;
	return rr;
}

int query_parse_question(void *message, void *data,
		char *name, __u16 *type, __u16 *class)
{
	int name_len, offset;

	name_len = query_parse_name(message, data, name, &offset);
	data += offset;

	*type = ntohs(((__u16 *)data)[0]);
	*class = ntohs(((__u16 *)data)[1]);

	uloga("%s: question: name: '%s', type: %d, class: %d.\n",
			query_spaces, name, *type, *class);

	return offset + 4;
}

int query_parse_answer(void *data)
{
	void *rrh;
	struct query_header *h = data;
	struct rr *rr;
	int i;
	unsigned int offset;
	char name[QUERY_RR_NAME_MAX_SIZE];
	__u16 type, class;

	query_parse_header(h);

	rrh = (void *)(h + 1);

	for (i=0; i<h->question_num; ++i) {
		rrh += query_parse_question(data, rrh, name, &type, &class);
	}

	for (i=0; i<h->answer_num + h->auth_num + h->addon_num; ++i) {
		offset = 0;
		rr = query_parse_rr(data, rrh, &offset);
		if (!rr)
			break;

		free(rr);

		rrh += offset;
	}

	if (!rr)
		return -EINVAL;

	return 0;
}

int query_add_rr_noname(void *data, struct rr *rr)
{
	__u16 *a = data;
	__u32 *ttl;

	a[0] = htons(rr->type);
	a[1] = htons(rr->class);
	ttl = (__u32 *)&a[2];
	ttl[0] = htonl(rr->ttl);
	a[4] = htons(rr->rdlen);
	memcpy(&a[5], rr->rdata, rr->rdlen);

	return rr->rdlen + 10;
}

int query_add_rr(void *answer, struct rr *rr)
{
	__u16 *a = answer;

	a[0] = htons(0xc000 | sizeof(struct query_header));

	return 2 + query_add_rr_noname(&a[1], rr);
}