/**************************************
                *     AUTHOR: Federico Tomassini        *
               *     Copyright (C) Federico Tomassini    *
              *     Contact effetom@gmail.com             *
             ***********************************************
             *******          BEGIN 3/2006          ********
*************************************************************************
*                                                                       *
*  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 "andns_lib.h"
#include "andns_net.h"
#include "log.h"
#include "err_errno.h"
#include "xmalloc.h"

#include <arpa/inet.h>
#include <zlib.h>

int
andns_compress(char *src, int srclen)
{
	int res;
	uLongf space;

	src += ANDNS_HDR_SZ;
	srclen -= ANDNS_HDR_SZ;
	space = compressBound(srclen);

	unsigned char dst[space + ANDNS_HDR_Z];

	/* The first four bytes will store
	 * the uncompressed size */
	res = compress2(dst + ANDNS_HDR_Z, &space, (u_char *) src, srclen,
					ANDNS_COMPR_LEVEL);
	if (res != Z_OK)
		err_ret(ERR_ZLIBCP, -1);
	if (space >= srclen - ANDNS_HDR_Z)	/* We have to consider the four
										   bytes too */
		err_ret(ERR_ZLIBNU, -1);	/* This is a
									   silent return */
	res = htonl(srclen);
	memcpy(dst, &res, ANDNS_HDR_Z);
	memcpy(src, dst, space);

	return (int) space;
}

char *
andns_uncompress(char *src, int srclen, int *dstlen)
{
	unsigned char *dst;
	uLongf space;
	int res;
	int c_len;
	const int hdrsz = ANDNS_HDR_SZ + ANDNS_HDR_Z;

	memcpy(&c_len, src + ANDNS_HDR_SZ, ANDNS_HDR_Z);
	c_len = ntohl(c_len);
	dst = xmalloc(c_len + ANDNS_HDR_SZ);

	space = c_len;

	res =
		uncompress(dst + ANDNS_HDR_SZ, &space, (u_char *) src + hdrsz,
				   srclen - hdrsz);
	if (res != Z_OK) {
		xfree(dst);
		err_ret(ERR_ZLIBUP, NULL);
	}
	if ((int) space != c_len) {
		xfree(dst);
		err_ret(ERR_ANDMAP, NULL);
	}

	memcpy(dst, src, ANDNS_HDR_SZ);
	*dstlen = c_len + ANDNS_HDR_SZ;

	return (char *) dst;
}

/*
 * Takes the buffer stream and translate headers to
 * andns_pkt struct.
 * Returns ALWAYS 4. The pkt_len has to be controlled
 * elsewhere.
 */
int
a_hdr_u(char *buf, andns_pkt * ap)
{
	uint8_t c;
	uint16_t s;
	char *start_buf;

	start_buf = buf;

	ap->r = *(buf + 1) & 0x01;
	memcpy(&s, buf, 2);
	ap->id = ntohs(s);
	ap->id >>= 1;
	buf += 2;

	memcpy(&c, buf, sizeof(uint8_t));
	ap->qr = (c >> 7) & 0x01;
	ap->p = c & 0x40 ? ANDNS_PROTO_UDP : ANDNS_PROTO_TCP;
	ap->z = c & 0x20;
	ap->qtype = (c >> 3) & 0x03;
	ap->ancount = (c << 1) & 0x0e;

	buf++;
	memcpy(&c, buf, sizeof(uint8_t));
	if (((*buf) & 0x80))
		ap->ancount++;

	ap->ipv = (c >> 6) & 0x01;
	ap->nk = (c >> 4) & 0x03;
	ap->rcode = c & 0x0f;
	return ANDNS_HDR_SZ;
}

/*
 * Translate the andns_pkt question stream to andns_pkt struct.
 * -1 on error. Bytes readed otherwise.
 *  NOTE: The qst-data size is controlled: apkt won't need
 *  this control.
 */
int
a_qst_u(char *buf, andns_pkt * ap, int limitlen)
{
	int ret;
	uint16_t s;

	if (limitlen < 3)
		err_ret(ERR_ANDMAP, -1);

	switch (ap->qtype) {
	case AT_A:
		memcpy(&s, buf, 2);
		ap->service = ntohs(s);
		buf += 2;
		if (ap->nk == NTK_REALM) {
			ap->qstlength = ANDNS_HASH_H;
			if (ap->qstlength > limitlen - 2)
				err_ret(ERR_ANDPLB, -1);
			AP_ALIGN(ap);
			memcpy(ap->qstdata, buf, ANDNS_HASH_H);
			ret = ANDNS_HASH_H + 2;
		} else if (ap->nk == INET_REALM) {
			memcpy(&s, buf, 2);
			ap->qstlength = ntohs(s);
			buf += 2;
			if (ap->qstlength >= ANDNS_MAX_QST_LEN ||
				ap->qstlength > limitlen - 4)
				err_ret(ERR_ANDPLB, -1);
			AP_ALIGN(ap);
			memcpy(ap->qstdata, buf, ap->qstlength);
			ret = ap->qstlength + 4;
		} else
			return -1;
		break;
	case AT_PTR:
		ap->qstlength = ap->ipv ? 16 : 4;
		if (ap->qstlength > limitlen)
			err_ret(ERR_ANDMAP, -1)
				AP_ALIGN(ap);
		memcpy(ap->qstdata, buf, ap->qstlength);
		ret = ap->qstlength;
		break;
	case AT_G:
		if (ap->nk != NTK_REALM)
			err_ret(ERR_ANDMAP, -1);
		ap->qstlength = ANDNS_HASH_H;
		if (ap->qstlength > limitlen)
			err_ret(ERR_ANDPLB, -1);
		AP_ALIGN(ap);
		memcpy(ap->qstdata, buf, ANDNS_HASH_H);
		ret = ap->qstlength;
		break;
	default:
		debug(DBG_INSANE, "In a_qst_u: unknow query type.");
		err_ret(ERR_ANDMAP, -1)
	}
	return ret;
}

int
a_answ_u(char *buf, andns_pkt * ap, int limitlen)
{
	uint16_t alen;
	andns_pkt_data *apd;
	int limit;

	if (limitlen < 3)
		err_ret(ERR_ANDMAP, -1);
	switch (ap->qtype) {
	case AT_A:
		limit = 2;
		if (limitlen < limit)
			err_ret(ERR_ANDPLB, -1);
		apd = andns_add_answ(ap);
		if (*buf & 0x40) {
			apd->m |= APD_IP;
			if (*buf & 0x80)
				apd->m |= APD_MAIN_IP;
			limit = ap->ipv ? 16 : 4;
		} else
			limit = ANDNS_HASH_H;
		if (limitlen < limit + 2)
			err_ret(ERR_ANDPLB, -1);
		apd->wg = (*buf & 0x3f);
		apd->prio = (*(buf + 1));
		apd->rdlength = limit;
		APD_ALIGN(apd);
		memcpy(apd->rdata, buf + 2, limit);
		limit += 2;
		break;
	case AT_PTR:
		memcpy(&alen, buf, 2);
		alen = ntohs(alen);
		if (alen + 2 > limitlen)
			err_ret(ERR_ANDPLB, -1);
		if (alen > ANDNS_MAX_DATA_LEN)
			err_ret(ERR_ANDPLB, -1);
		apd = andns_add_answ(ap);
		apd->rdlength = alen;
		APD_ALIGN(apd);
		memcpy(apd->rdata, buf + 2, alen);
		limit = alen + 2;
		break;
	case AT_G:
		if (limitlen < 8)
			err_ret(ERR_ANDMAP, -1);
		apd = andns_add_answ(ap);
		if (*buf & 0x40) {
			apd->m |= APD_IP;
			if (*buf & 0x80)
				apd->m |= APD_MAIN_IP;
		}
		apd->m |= *buf & 0x20 ? APD_UDP : APD_TCP;
		apd->wg = (*buf & 0x1f);
		apd->prio = (*(buf + 1));
		buf += 2;
		memcpy(&alen, buf, 2);
		apd->service = ntohs(alen);
		buf += 2;
		if (apd->m & APD_IP)
			apd->rdlength = (ap->ipv ? 16 : 4);
		else
			apd->rdlength = ANDNS_HASH_H;
		limit = 4 + apd->rdlength;
		if (limitlen < limit)
			err_ret(ERR_ANDPLB, -1);
		APD_ALIGN(apd);
		memcpy(apd->rdata, buf, apd->rdlength);
		break;
	default:
		err_ret(ERR_ANDMAP, -1);
	}
	return limit;
}

int
a_answs_u(char *buf, andns_pkt * ap, int limitlen)
{
	int ancount, i;
	int offset = 0, res;
	uint16_t alen;

	if (ap->qtype == AT_G) {
		memcpy(&alen, buf, sizeof(uint16_t));
		ap->ancount = ntohs(alen);
		offset += 2;
	}
	ancount = ap->ancount;
	for (i = 0; i < ancount; i++) {
		res = a_answ_u(buf + offset, ap, limitlen - offset);
		if (res == -1) {
			error(err_str);
			err_ret(ERR_ANDMAD, -1);
		}
		offset += res;
	}
	return offset;
}

/*
 * This is a main function: takes the pkt-buf and translate
 * it in structured data.
 * It cares about andns_pkt allocation.
 * The apkt is allocate here.
 *
 * Returns:
 * -1 on E_INTRPRT
 *  0 if pkt must be discarded.
 *  Number of bytes readed otherwise
 */
int
a_u(char *buf, int pktlen, andns_pkt ** app)
{
	andns_pkt *ap;
	int offset, res;
	int limitlen, u_len;
	char *u_buf;

	if (pktlen < ANDNS_HDR_SZ)
		err_ret(ERR_ANDPLB, 0);
	*app = ap = create_andns_pkt();
	offset = a_hdr_u(buf, ap);

	if (ap->z) {				/* Controls the space to read
								   uncompressed size */
		if (pktlen < ANDNS_HDR_SZ + ANDNS_HDR_Z) {
			destroy_andns_pkt(ap);
			err_ret(ERR_ANDPLB, 0);
		}
		if (!(u_buf = andns_uncompress(buf, pktlen, &u_len)))
			goto andmap;
		destroy_andns_pkt(ap);
		ANDNS_UNSET_Z(u_buf);
		res = a_u(u_buf, u_len, app);
		xfree(u_buf);
		return res;
	}
	buf += offset;
	limitlen = pktlen - offset;
	if ((res = a_qst_u(buf, ap, limitlen)) == -1)
		goto andmap;
	offset += res;
	if (!ap->ancount)			/*No answers */
		return offset;
	buf += res;
	limitlen -= res;
	if ((res = a_answs_u(buf, ap, limitlen)) == -1)
		goto andmap;
	offset += res;
	if (offset != pktlen)
		error
			("In a_u(): pktlen differs from readed contents: ID query %d.",
			 ap->id);
	return offset;
  andmap:
	destroy_andns_pkt(ap);
	error(err_str);
	err_ret(ERR_ANDMAP, -1);
}

int
a_hdr_p(andns_pkt * ap, char *buf)
{
	uint16_t s;
	uint8_t an;

	ap->id <<= 1;
	s = htons(ap->id);
	memcpy(buf, &s, sizeof(uint16_t));
	if (ap->r)
		*(buf + 1) |= 0x01;
	else
		*(buf + 1) &= 0xfe;
	buf += 2;
	if (ap->qr)
		(*buf) |= 0x80;
	if (ap->p)
		(*buf) |= 0x40;
	if (ap->z)
		(*buf) |= 0x20;
	(*buf) |= ((ap->qtype) << 3);
	an = ap->ancount;
	(*buf++) |= ((an) >> 1);
	(*buf) |= ((ap->ancount) << 7);
	if (ap->ipv)
		*buf |= 0x40;
	(*buf) |= ((ap->nk) << 4);
	(*buf) |= (ap->rcode);
	return ANDNS_HDR_SZ;
}

int
a_qst_p(andns_pkt * ap, char *buf, int limitlen)
{
	int ret = 0;
	uint16_t s;
	int limit;

	switch (ap->qtype) {
	case AT_A:
		limit = ap->nk == NTK_REALM ? ANDNS_HASH_H + 2 : ap->qstlength + 4;
		if (limitlen < limit)
			err_ret(ERR_ANDMAD, -1);
		s = htons(ap->service);
		memcpy(buf, &s, 2);
		buf += 2;				/* here INET and NTK REALM change */
		if (ap->nk == NTK_REALM) {
			memcpy(buf, ap->qstdata, ANDNS_HASH_H);
			ret = ANDNS_HASH_H + 2;
		} else if (ap->nk == INET_REALM) {
			s = htons(ap->qstlength);
			memcpy(buf, &s, 2);
			buf += 2;
			memcpy(buf, ap->qstdata, ap->qstlength);
			ret = ap->qstlength + 4;
		} else
			err_ret(ERR_ANDMAD, -1);
		break;
	case AT_PTR:
		limit = ap->ipv ? 16 : 4;
		if (limitlen < limit)
			err_ret(ERR_ANDMAD, -1);
		memcpy(buf, ap->qstdata, limit);
		ret = limit;
		break;
	case AT_G:
		limit = ANDNS_HASH_H;
		if (limitlen < limit)
			err_ret(ERR_ANDMAD, -1);
		memcpy(buf, ap->qstdata, ANDNS_HASH_H);
		ret = ANDNS_HASH_H;
		break;
	default:
		debug(DBG_INSANE, "In a_qst_p: unknow query type.");
		err_ret(ERR_ANDMAD, -1);
		break;
	}
	return ret;
}

int
a_answ_p(andns_pkt * ap, andns_pkt_data * apd, char *buf, int limitlen)
{
	uint16_t s;
	int limit;
	int ret;

	switch (ap->qtype) {
	case AT_A:
		limit = ap->ipv ? 16 : 4;
		if (limitlen < limit + 2)
			err_ret(ERR_ANDPLB, -1);
		if (apd->m)
			*buf |= 0x80;
		*buf++ |= (apd->wg & 0x7f);
		*buf++ |= apd->prio;
		memcpy(buf, apd->rdata, limit);
		ret = limit + 2;
		break;
	case AT_PTR:
		if (limitlen < apd->rdlength + 2)
			err_ret(ERR_ANDPLB, -1);
		s = htons(apd->rdlength);
		memcpy(buf, &s, sizeof(uint16_t));
		buf += 2;
		memcpy(buf, apd->rdata, apd->rdlength);
		ret = apd->rdlength + 2;
		break;
	case AT_G:					/* TODO */
		if (limitlen < 4)
			err_ret(ERR_ANDPLB, -1);
		if (apd->m == 1)
			(*buf) |= 0xc0;
		else if (apd->m)
			(*buf) |= 0x40;
		*buf++ |= (apd->wg & 0x3f);
		*buf++ |= apd->prio;
		s = htons(apd->service);
		memcpy(buf, &s, 2);
		if (apd->m) {
			limit = ap->ipv ? 16 : 4;
			if (limitlen < limit + 4)
				err_ret(ERR_ANDPLB, -1);
			memcpy(buf, apd->rdata, limit);
			ret = limit + 4;
		} else {
			limit = strlen(apd->rdata);
			if (limitlen < limit + 6)
				err_ret(ERR_ANDPLB, -1);
			s = htons(limit);
			memcpy(buf, &s, 2);
			buf += 2;
			memcpy(buf, apd->rdata, limit);
			ret = limit + 6;
		}
		break;
	default:
		debug(DBG_INSANE, "In a_answ_p(): unknow query type.");
		err_ret(ERR_ANDMAD, -1);
		break;
	}
	return ret;
}

int
a_answs_p(andns_pkt * ap, char *buf, int limitlen)
{
	andns_pkt_data *apd;
	int i;
	int offset = 0, res;
	uint16_t s;

	if (ap->qtype == AT_G) {
		if (limitlen < 2)
			err_ret(ERR_ANDPLB, -1);
		s = htons(ap->ancount);
		memcpy(buf, &s, 2);
		offset += 2;
	}
	apd = ap->pkt_answ;
	for (i = 0; i < ap->ancount && apd; i++) {
		if ((res =
			 a_answ_p(ap, apd, buf + offset, limitlen - offset)) == -1) {
			error(err_str);
			err_ret(ERR_ANDMAD, -1);
		}
		offset += res;
		apd = apd->next;
	}
	return offset;
}

int
a_p(andns_pkt * ap, char *buf)
{
	int offset, res;

	memset(buf, 0, ANDNS_MAX_SZ);

	offset = a_hdr_p(ap, buf);
	buf += offset;
	if ((res = a_qst_p(ap, buf, ANDNS_MAX_SZ - offset)) == -1)
		goto server_fail;
	offset += res;
	buf += res;
	if (ap->ancount) {
		if ((res = a_answs_p(ap, buf, ANDNS_MAX_SZ - offset)) == -1)
			goto server_fail;
		offset += res;
	}
	destroy_andns_pkt(ap);
	/* Compression */
	if (offset > ANDNS_COMPR_THRESHOLD) {
		res = andns_compress(buf, offset);
		if (res == -1)
			error(err_str);
		else
			return res;
	}
	/* end compression */
	return offset;
  server_fail:
	destroy_andns_pkt(ap);
	error(err_str);
	err_ret(ERR_ANDMAD, -1);
}


/* MEM */

	/* Borning functions */
andns_pkt *
create_andns_pkt(void)
{
	andns_pkt *ap;
	ap = xmalloc(ANDNS_PKT_SZ);
	memset(ap, 0, ANDNS_PKT_SZ);
	return ap;
}

andns_pkt_data *
create_andns_pkt_data(void)
{
	andns_pkt_data *apd;
	apd = xmalloc(ANDNS_PKT_DATA_SZ);
	memset(apd, 0, ANDNS_PKT_DATA_SZ);
	return apd;
}

andns_pkt_data *
andns_add_answ(andns_pkt * ap)
{
	andns_pkt_data *apd, *a;

	apd = create_andns_pkt_data();
	a = ap->pkt_answ;
	if (!a) {
		ap->pkt_answ = apd;
		return apd;
	}
	while (a->next)
		a = a->next;
	a->next = apd;
	return apd;
}

	/* Death functions */
void
destroy_andns_pkt_data(andns_pkt_data * apd)
{
	if (apd->rdata)
		xfree(apd->rdata);
	xfree(apd);
}

void
andns_del_answ(andns_pkt * ap)
{
	andns_pkt_data *apd, *apdt;

	apd = ap->pkt_answ;
	if (!apd)
		return;
	apdt = apd->next;
	while (apdt) {
		apd = apdt;
		apdt = apdt->next;
	}
	apd->next = NULL;
	destroy_andns_pkt_data(apdt);
}

void
destroy_andns_pkt_datas(andns_pkt * ap)
{
	andns_pkt_data *apd, *apd_t;
	apd = ap->pkt_answ;
	while (apd) {
		apd_t = apd->next;
		destroy_andns_pkt_data(apd);
		apd = apd_t;
	}
}

void
destroy_andns_pkt(andns_pkt * ap)
{
	if (ap->qstdata)
		xfree(ap->qstdata);
	destroy_andns_pkt_datas(ap);
	xfree(ap);
}