/*---------------------------------------------------------------------------*/
/* dns_query.cpp                                                             */
/* copyright (c) innovaphone 2024                                            */
/* DNS query response parser                                                 */
/*---------------------------------------------------------------------------*/

#include <platform.h>
#include "common/ilib/str.h"
#include "sdk/common/interface/dns.h"

#include "sdk/common/lib/dns_query.h"

#define IS_QR_BIT(flags)  (byte)(flags&0x8000 ? true : false)

#define DNS_MAX_LABEL   63
#define DNS_MAX_NAME    255

struct DnsQueryHeader {
    word transationId;
    word flags;
    word questions;
    word answers;
    word authorities;
    word additionals;
};


/*--------------------------------------------------------------------------------------------*/
/*                               IDnsQueryResult::Parse                                     */
/*--------------------------------------------------------------------------------------------*/

class IDnsQueryResult * IDnsQueryResult::Parse(const byte * answer, unsigned length)
{
    if (answer == nullptr || length == 0 || length < sizeof(DnsQueryHeader)) return nullptr;
    
	DnsQueryResponse * response = new class DnsQueryResponse(answer, length);
	if (response->GetError()) {
		delete response;
		response = nullptr;
	}

    return response;
}


/*--------------------------------------------------------------------------------------------*/
/*                               class DnsQueryResponse                                       */
/*--------------------------------------------------------------------------------------------*/

DnsQueryResponse::DnsQueryResponse(const byte * answer, unsigned length) {
	this->error = true;

	DnsQueryHeader header;
	const byte * msg = answer;
	const byte * msgEnd = answer + length;

	header.transationId = READ_WORD_NET(&answer[0]);
	header.flags = READ_WORD_NET(&answer[2]);
	header.questions = READ_WORD_NET(&answer[4]);
	header.answers = READ_WORD_NET(&answer[6]);
	header.authorities = READ_WORD_NET(&answer[8]);
	header.additionals = READ_WORD_NET(&answer[10]);

	if (!IS_QR_BIT(header.flags)) return;

	msg += sizeof(DnsQueryHeader);

	// parse the questions first -> ignore the content, not needed for our purpose
	char dnsName[DNS_MAX_NAME];
	word queryType, queryClass;

	for (int i = 0; i < header.questions; i++) {
		msg = this->ParseQuestion(answer, msg, msgEnd, dnsName, dnsName + sizeof(dnsName), &queryType, &queryClass);
		if (msg == nullptr) return;
	}

	msg = this->ParseSection(answer, msg, msgEnd, header.answers, &this->answers);
	if (msg == nullptr) return;
	msg = this->ParseSection(answer, msg, msgEnd, header.authorities, &this->authorities);
	if (msg == nullptr) return;
	msg = this->ParseSection(answer, msg, msgEnd, header.additionals, &this->additionals);
	if (msg == nullptr) return;

	this->error = false;	
}

DnsQueryResponse::~DnsQueryResponse()
{
	for (auto entry : this->answers) delete entry;
	for (auto entry : this->authorities) delete entry;
	for (auto entry : this->additionals) delete entry;
}

const byte * DnsQueryResponse::ParseQuestion(const byte * msg, const byte * start, const byte * end, char * name, char * nameEnd, word * queryType, word * queryClass)
{
	if (!start || !end || !name || !nameEnd || !queryType || !queryClass) return nullptr;
	*name = '\0';
	start = this->ParseName(msg, start, end, name, nameEnd);
	if (!start || start >= end) return nullptr;

	*queryType = READ_WORD_NET(start); start += sizeof(word);
	if (start >= end) return nullptr;
	*queryClass = READ_WORD_NET(start); start += sizeof(word);
	if (start > end) return nullptr;

	return start;
}

const byte * DnsQueryResponse::ParseName(const byte * msg, const byte * start, const byte * end, char * name, char * nameEnd)
{
	if (!start) return nullptr;

	if (*start & 0xC0) {
		/* pointer field */
		if (start + sizeof(word) > end) return nullptr;
		word offset = READ_WORD_NET(start) & 0x3FFF; start += sizeof(word);

		if (msg + offset >= end || !this->ParseName(msg, &msg[offset], end, name, nameEnd)) return nullptr;
	}
	else {
		/* length field */
		if (start + sizeof(byte) > end) return nullptr;
		int len = *start++;

		if (!len) {
			/* could be root */
			*name++ = '\0';
			return(start);
		}

		if (len > DNS_MAX_LABEL
			|| start + len > end
			|| name + len + sizeof(byte) > nameEnd) return nullptr;

		while (len--) *name++ = *start++;
		if (*start) {
			*name++ = '.';    /* a label follows */
			start = this->ParseName(msg, start, end, name, nameEnd);
		}
		else {
			*name++ = '\0';   /* reached end */
			start++;          /* skip zero length byte */
		}
	}

	return start;
}

const byte * DnsQueryResponse::ParseSection(const byte * msg, const byte * start, const byte * end, word count, std::vector<class IDnsQueryEntry *> * list)
{
	if (!start) return nullptr;

	// each record inside a section looks like this:

/*	+-- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- +
	/ NAME																		   /
	+-- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- +
	| TYPE	2 Byte																   |
	+-- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- +
	| CLASS	2 Byte																   |
	+-- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- +
	| TTL	4 Byte																   |
	+-- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- +
	| RDLENGTH																	   |
	+-- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- +
	/ RDATA		various according to it's TYPE		    					       /
	+-- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + -- + */

	for (word i = 0; i < count; i++) {
		char name[DNS_MAX_NAME];
		word queryType;
		word queryClass;
		dword ttl;
		word rdLength;

		start = this->ParseName(msg, start, end, name, name + sizeof(name));
		if (start == nullptr) return nullptr;

		if (start + sizeof(word) > end) return nullptr;
		queryType = READ_WORD_NET(start); start += sizeof(word);
		if (start + sizeof(word) > end) return nullptr;
		queryClass = READ_WORD_NET(start); start += sizeof(word);
		if (start + sizeof(dword) > end) return nullptr;
		ttl = READ_DWORD_NET(start); start += sizeof(dword);
		if (start + sizeof(word) > end) return nullptr;
		rdLength = READ_WORD_NET(start); start += sizeof(word);

		if (queryClass != 1	/* the arpa internet */) {
			start += rdLength;
			continue;
		}

		class IDnsQueryEntry * entry = nullptr;

		switch (queryType)
		{
		case DNS_QUERY_TYPE_A:		entry = this->ParseA(msg, start, end, name, ttl); break;
		case DNS_QUERY_TYPE_AAAA:	entry = this->ParseAAAA(msg, start, end, name, ttl); break;
		case DNS_QUERY_TYPE_CNAME:	entry = this->ParseCNAME(msg, start, end, name, ttl); break;
		case DNS_QUERY_TYPE_MX:		entry = this->ParseMX(msg, start, end, name, ttl); break;
		case DNS_QUERY_TYPE_NS:		entry = this->ParseNS(msg, start, end, name, ttl); break;
		case DNS_QUERY_TYPE_SRV:	entry = this->ParseSRV(msg, start, end, name, ttl); break;
		case DNS_QUERY_TYPE_TXT:	entry = this->ParseTXT(msg, start, end, rdLength, name, ttl); break;
		case DNS_QUERY_TYPE_SOA:	entry = this->ParseSOA(msg, start, end, name, ttl); break;
		case DNS_QUERY_TYPE_PTR:	entry = this->ParsePTR(msg, start, end, name, ttl); break;
		case DNS_QUERY_TYPE_NAPTR:	entry = this->ParseNAPTR(msg, start, end, name, ttl); break;
		default:
			// unsupported record type
			start += rdLength;
			continue;
		}
		start += rdLength;
		if (entry == nullptr) return nullptr;
		else list->push_back(entry);
	}

	return start;
}

class IDnsQueryEntry * DnsQueryResponse::ParseA(const byte * msg, const byte * start, const byte * end, const char * name, dword ttl)
{	
	if (start + sizeof(dword) > end) return nullptr;

	char addr[32];
	// no str::from_ip here, as I didn't find a simple way to copy 4 bytes into IPaddr so that IPaddr is not recognized as IP6 address ...
	_snprintf(addr, sizeof(addr), "%u.%u.%u.%u", start[0], start[1], start[2], start[3]);

	return new class DnsQueryEntryA(name, ttl, addr);
}

class IDnsQueryEntry * DnsQueryResponse::ParseAAAA(const byte * msg, const byte * start, const byte * end, const char * name, dword ttl)
{
	IPaddr addr = ip_anyaddr;

	if (start + sizeof(dword) * 4 > end) return nullptr;

	char ipAddr[32];
	memcpy(&addr.b[0], start, sizeof(dword) * 4);
	str::from_ip(ipAddr, sizeof(ipAddr), &addr);

	return new class DnsQueryEntryAAAA(name, ttl, ipAddr);
}

class IDnsQueryEntry * DnsQueryResponse::ParseCNAME(const byte * msg, const byte * start, const byte * end, const char * name, dword ttl)
{
	char canonicalName[DNS_MAX_NAME];
	if (this->ParseName(msg, start, end, canonicalName, canonicalName + sizeof(canonicalName)) == nullptr) return nullptr;

	return new class DnsQueryEntryCNAME(name, ttl, canonicalName);
}

class IDnsQueryEntry * DnsQueryResponse::ParseMX(const byte * msg, const byte * start, const byte * end, const char * name, dword ttl)
{
	word preference;
	char mailServer[DNS_MAX_NAME];

	if (start + sizeof(word) > end) return nullptr;
	preference = READ_WORD_NET(start); start += sizeof(word);
	if (this->ParseName(msg, start, end, mailServer, mailServer + sizeof(mailServer)) == nullptr) return nullptr;

	return new class DnsQueryEntryMX(name, ttl, preference, mailServer);
}

class IDnsQueryEntry * DnsQueryResponse::ParseNS(const byte * msg, const byte * start, const byte * end, const char * name, dword ttl)
{
	char nameServer[DNS_MAX_NAME];
	if (this->ParseName(msg, start, end, nameServer, nameServer + sizeof(nameServer)) == nullptr) return nullptr;

	return new class DnsQueryEntryNS(name, ttl, nameServer);
}

class IDnsQueryEntry * DnsQueryResponse::ParseSRV(const byte * msg, const byte * start, const byte * end, const char * name, dword ttl)
{
	word priority;
	word weight;
	word port;
	char target[DNS_MAX_NAME];

	if (start + 3 * sizeof(word) > end) return nullptr;
	priority = READ_WORD_NET(start); start += sizeof(word);
	weight = READ_WORD_NET(start); start += sizeof(word);
	port = READ_WORD_NET(start); start += sizeof(word);
	if (this->ParseName(msg, start, end, target, target + sizeof(target)) == nullptr) return nullptr;

	return new class DnsQueryEntrySRV(name, ttl, priority, weight, port, target);
}

class IDnsQueryEntry * DnsQueryResponse::ParseTXT(const byte * msg, const byte * start, const byte * end, word rdLength, const char * name, dword ttl)
{
	class DnsQueryEntryTXT * txt = new class DnsQueryEntryTXT(name, ttl);
	while (rdLength > 0) {
		if (start + sizeof(byte) > end) {
			delete txt;
			return nullptr;
		}
		byte length = start[0]; start++; rdLength--;
		if (start + length > end || length > rdLength) {
			delete txt;
			return nullptr;
		}
		txt->AddText(length, start);
		rdLength -= length;	// length of text + length byte itself
		start += length;
	}
	return txt;
}

class IDnsQueryEntry * DnsQueryResponse::ParseSOA(const byte * msg, const byte * start, const byte * end, const char * name, dword ttl)
{
	char primaryNS[DNS_MAX_NAME];
	char responsibleMailbox[DNS_MAX_NAME];
	dword serialNumber;
	dword refreshInterval;
	dword retryInterval;
	dword expireLimit;
	dword minimumTTL;

	start = this->ParseName(msg, start, end, primaryNS, primaryNS + sizeof(primaryNS));
	if (start == nullptr) return nullptr; 
	start = this->ParseName(msg, start, end, responsibleMailbox, responsibleMailbox + sizeof(responsibleMailbox));
	if (start == nullptr) return nullptr;
	if (start + sizeof(dword) > end) return nullptr;
	serialNumber = READ_DWORD_NET(start); start += sizeof(dword);
	if (start + sizeof(dword) > end) return nullptr;
	refreshInterval = READ_DWORD_NET(start); start += sizeof(dword);
	if (start + sizeof(dword) > end) return nullptr;
	retryInterval = READ_DWORD_NET(start); start += sizeof(dword);
	if (start + sizeof(dword) > end) return nullptr;
	expireLimit = READ_DWORD_NET(start); start += sizeof(dword);
	if (start + sizeof(dword) > end) return nullptr;
	minimumTTL = READ_DWORD_NET(start); start += sizeof(dword);

	return new class DnsQueryEntrySOA(name, ttl, primaryNS, responsibleMailbox, serialNumber, refreshInterval, retryInterval, expireLimit, minimumTTL);
}

class IDnsQueryEntry * DnsQueryResponse::ParsePTR(const byte * msg, const byte * start, const byte * end, const char * name, dword ttl)
{
	char domainName[DNS_MAX_NAME];
	start = this->ParseName(msg, start, end, domainName, domainName + sizeof(domainName));
	if (start == nullptr) return nullptr;

	return new class DnsQueryEntryPTR(name, ttl, domainName);
}

class IDnsQueryEntry * DnsQueryResponse::ParseNAPTR(const byte * msg, const byte * start, const byte * end, const char * name, dword ttl)
{
	word order;
	word preference;
	byte len;
	char flags[DNS_MAX_NAME + 1];
	char services[DNS_MAX_NAME + 1];
	char regex[DNS_MAX_NAME + 1];
	char replacement[DNS_MAX_NAME];

	if (start + sizeof(word) > end) return nullptr;
	order = READ_WORD_NET(start); start += sizeof(word);
	if (start + sizeof(word) > end) return nullptr;
	preference = READ_WORD_NET(start); start += sizeof(word);
	if (start + sizeof(byte) > end) return nullptr;
	len = *start; start++;
	if (start + len > end) return nullptr;
	memcpy(flags, start, len);
	flags[len] = '\0';
	if (start + sizeof(byte) > end) return nullptr;
	len = *start; start++;
	if (start + len > end) return nullptr;
	memcpy(services, start, len);
	services[len] = '\0';
	if (start + sizeof(byte) > end) return nullptr;
	len = *start; start++;
	if (start + len > end) return nullptr;
	memcpy(regex, start, len);
	regex[len] = '\0';
	start = this->ParseName(msg, start, end, replacement, replacement + sizeof(replacement));
	if (start == nullptr) return nullptr;

	return new class DnsQueryEntryNAPTR(name, ttl, order, preference, flags, services, regex, replacement);
}

/*--------------------------------------------------------------------------------------------*/
/*                               class DnsQueryEntry                                          */
/*--------------------------------------------------------------------------------------------*/

DnsQueryEntry::DnsQueryEntry(const char * name, dns_query_type_sdk type, dword ttl)
{
	this->name = _strdup(name);
	this->type = type;
	this->ttl = ttl;
}

DnsQueryEntry::~DnsQueryEntry()
{
	free(this->name);
}


/*--------------------------------------------------------------------------------------------*/
/*                               class DnsQueryEntryA                                         */
/*--------------------------------------------------------------------------------------------*/

DnsQueryEntryA::DnsQueryEntryA(const char * name, dword ttl, const char * address) :
	DnsQueryEntry(name, DNS_QUERY_TYPE_A, ttl)
{
	this->address = _strdup(address);
}

DnsQueryEntryA::~DnsQueryEntryA()
{
	free(this->address);
}


/*--------------------------------------------------------------------------------------------*/
/*                               class DnsQueryEntryAAAA                                      */
/*--------------------------------------------------------------------------------------------*/

DnsQueryEntryAAAA::DnsQueryEntryAAAA(const char * name, dword ttl, const char * address) :
	DnsQueryEntry(name, DNS_QUERY_TYPE_AAAA, ttl)
{
	this->address = _strdup(address);
}

DnsQueryEntryAAAA::~DnsQueryEntryAAAA()
{
	free(this->address);
}


/*--------------------------------------------------------------------------------------------*/
/*                               class DnsQueryEntryCNAME                                     */
/*--------------------------------------------------------------------------------------------*/

DnsQueryEntryCNAME::DnsQueryEntryCNAME(const char * name, dword ttl, const char * canonicalName) :
	DnsQueryEntry(name, DNS_QUERY_TYPE_CNAME, ttl)
{
	this->canonicalName = _strdup(canonicalName);
}

DnsQueryEntryCNAME::~DnsQueryEntryCNAME()
{
	free(this->canonicalName);
}


/*--------------------------------------------------------------------------------------------*/
/*                               class DnsQueryEntryMX                                        */
/*--------------------------------------------------------------------------------------------*/

DnsQueryEntryMX::DnsQueryEntryMX(const char * name, dword ttl, word preference, const char * mailServer) :
	DnsQueryEntry(name, DNS_QUERY_TYPE_MX, ttl)
{
	this->preference = preference;
	this->mailServer = _strdup(mailServer);
}

DnsQueryEntryMX::~DnsQueryEntryMX()
{
	free(this->mailServer);
}


/*--------------------------------------------------------------------------------------------*/
/*                               class DnsQueryEntryNS                                        */
/*--------------------------------------------------------------------------------------------*/

DnsQueryEntryNS::DnsQueryEntryNS(const char * name, dword ttl, const char * nameServer) :
	DnsQueryEntry(name, DNS_QUERY_TYPE_NS, ttl)
{
	this->nameServer = _strdup(nameServer);
}

DnsQueryEntryNS::~DnsQueryEntryNS()
{
	free(this->nameServer);
}


/*--------------------------------------------------------------------------------------------*/
/*                               class DnsQueryEntrySRV                                       */
/*--------------------------------------------------------------------------------------------*/

DnsQueryEntrySRV::DnsQueryEntrySRV(const char * name, dword ttl, word priority, word weight, word port, const char * target) :
	DnsQueryEntry(name, DNS_QUERY_TYPE_SRV, ttl)
{
	this->priority = priority;
	this->weight = weight;
	this->port = port;
	this->target = _strdup(target);
}

DnsQueryEntrySRV::~DnsQueryEntrySRV()
{
	free(this->target);
}


/*--------------------------------------------------------------------------------------------*/
/*                               class DnsQueryEntryTXT                                       */
/*--------------------------------------------------------------------------------------------*/

DnsQueryEntryTXT::DnsQueryEntryTXT(const char * name, dword ttl) :
	DnsQueryEntry(name, DNS_QUERY_TYPE_TXT, ttl)
{

}

DnsQueryEntryTXT::~DnsQueryEntryTXT()
{
	for (auto text : this->texts) free(text);
}

void DnsQueryEntryTXT::AddText(byte length, const byte * start)
{
	char * text = (char *)malloc(length + 1);
	memcpy(text, start, length);
	text[length] = '\0';
	this->texts.push_back(text);
}


/*--------------------------------------------------------------------------------------------*/
/*                               class DnsQueryEntrySOA                                       */
/*--------------------------------------------------------------------------------------------*/

DnsQueryEntrySOA::DnsQueryEntrySOA(const char * name, dword ttl, char * primaryNS, char * responsibleMailbox, dword serialNumber, dword refreshInterval, dword retryInterval, dword expireLimit, dword minimumTTL) :
	DnsQueryEntry(name, DNS_QUERY_TYPE_SOA, ttl)
{
	this->primaryNS = _strdup(primaryNS);
	this->responsibleMailbox = _strdup(responsibleMailbox);
	this->serialNumber = serialNumber;
	this->refreshInterval = refreshInterval;
	this->retryInterval = retryInterval;
	this->expireLimit = expireLimit;
	this->minimumTTL = minimumTTL;
}

DnsQueryEntrySOA::~DnsQueryEntrySOA()
{
	free(this->primaryNS);
	free(this->responsibleMailbox);
}


/*--------------------------------------------------------------------------------------------*/
/*                               class DnsQueryEntrySOA                                       */
/*--------------------------------------------------------------------------------------------*/

DnsQueryEntryPTR::DnsQueryEntryPTR(const char * name, dword ttl, const char * domainName) :
	DnsQueryEntry(name, DNS_QUERY_TYPE_PTR, ttl)
{
	this->domainName = _strdup(domainName);
}

DnsQueryEntryPTR::~DnsQueryEntryPTR()
{
	free(this->domainName);
}


/*--------------------------------------------------------------------------------------------*/
/*                               class DnsQueryEntrySOA                                       */
/*--------------------------------------------------------------------------------------------*/

DnsQueryEntryNAPTR::DnsQueryEntryNAPTR(const char * name, dword ttl, word order, word preference, const char * flags, const char * services, const char * regex, const char * replacement) :
	DnsQueryEntry(name, DNS_QUERY_TYPE_NAPTR, ttl)
{
	this->order = order;
	this->preference = preference;
	this->flags = _strdup(flags);
	this->services = _strdup(services);
	this->regex = _strdup(regex);
	this->replacement = _strdup(replacement);
}

DnsQueryEntryNAPTR::~DnsQueryEntryNAPTR()
{
	free(this->flags);
	free(this->services);
	free(this->regex);
	free(this->replacement);
}
