//..............................................................................
//
// Sniffer log DNS packets handling
//

import "io_Dns.jnc"
import "std_String.jnc"

//..............................................................................

class DnsProtoParser: AppProtoParser
{
protected:
	enum
	{
		PreviewLineCount = 2,
	}

protected:
	char const* m_name;

public:
	construct(char const* name = "DNS")
	{
		m_name = name;
	}

	override void representDigest(
		log.Representation* representation,
		char const* srcString,
		char const* dstString,
		void const* p,
		size_t size,
		int payloadPreviewIconIdx
		);

	override void representDetails(
		log.Representation* representation,
		std.Array* protoRangeArray,
		void const* p,
		size_t size,
		size_t baseOffset,
		uint_t foldFlags
		);

protected:
	virtual decodeName(std.StringBuilder* string)
	{
	}

	void const* formatName(
		std.StringBuilder* string,
		void const* p,
		void const* end,
		void const* base
		);

	void const* formatNameLink(
		std.StringBuilder* string,
		void const* p,
		void const* end,
		void const* base,
		size_t baseOffset
		);

	char const* getAnswerRrValueString(
		std.StringBuilder* string,
		io.DnsType type,
		void const* p,
		void const* end,
		void const* base
		);

	void const* representQuestionRrTableDigest(
		log.Representation* representation,
		std.StringBuilder* string,
		size_t count,
		void const* p,
		void const* end,
		void const* base
		);

	void const* representAnswerRrTableDigest(
		log.Representation* representation,
		std.StringBuilder* string,
		size_t questionCount,
		size_t answerCount,
		void const* p,
		void const* end,
		void const* base
		);

	void const* representQuestionRrTableDetails(
		log.Representation* representation,
		std.StringBuilder* string,
		char const* name,
		size_t count,
		void const* p,
		void const* end,
		void const* base,
		size_t baseOffset
		);

	void const* representAnswerRrTableDetails(
		log.Representation* representation,
		std.StringBuilder* string,
		char const* name,
		size_t count,
		void const* p,
		void const* end,
		void const* base,
		size_t baseOffset
		);

	void const* representQuestionRrDigest(
		log.Representation* representation,
		std.StringBuilder* string,
		void const* p,
		void const* end,
		void const* base
		);

	void const* representAnswerRrDigest(
		log.Representation* representation,
		std.StringBuilder* string,
		void const* p,
		void const* end,
		void const* base
		);

	void const* representQuestionRrDetails(
		log.Representation* representation,
		std.StringBuilder* string,
		size_t i,
		void const* p,
		void const* end,
		void const* base,
		size_t baseOffset
		);

	void const* representAnswerRrDetails(
		log.Representation* representation,
		std.StringBuilder* string,
		size_t i,
		void const* p,
		void const* end,
		void const* base,
		size_t baseOffset
		);
}

// . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

void DnsProtoParser.representDigest(
	log.Representation* representation,
	char const* srcString,
	char const* dstString,
	void const* p,
	size_t size,
	int payloadPreviewIconIdx
	)
{
	representSimpleProtoDigestNoPayload(
		representation,
		typeof(io.DnsHdr),
		m_name,
		srcString,
		dstString,
		p,
		size
		);

	if (size <= sizeof(io.DnsHdr))
		return;

	io.DnsHdr const* hdr = (io.DnsHdr const*)p;
	void const* end = p + size;
	p += sizeof(io.DnsHdr);
	size -= sizeof(io.DnsHdr);

	size_t lineCount = 0;
	std.StringBuilder string;

	if (hdr.m_questionRrCount)
		p = representQuestionRrTableDigest(
			representation,
			string,
			hdr.m_questionRrCount,
			p,
			end,
			hdr
			);

	size_t answerCount = hdr.m_answerRrCount + hdr.m_authorityRrCount + hdr.m_additionalRrCount;
	if (answerCount)
		p = representAnswerRrTableDigest(
			representation,
			string,
			hdr.m_questionRrCount,
			answerCount,
			p,
			end,
			hdr
			);
}

void DnsProtoParser.representDetails(
	log.Representation* representation,
	std.Array* protoRangeArray,
	void const* p,
	size_t size,
	size_t baseOffset,
	uint_t foldFlags
	)
{
	representation.m_lineAttr.m_backColor = ProtoColor.AppLayer;

	if (size < sizeof(io.DnsHdr))
	{
		representation.addHyperText("\e[1mDNS\e[m [INCOMPLETE]");
		return;
	}

	log.representStruct(
		representation,
		typeof(io.DnsHdr),
		m_name,
		NetSnifferLogPartCode.RawPacket,
		p,
		baseOffset,
		foldFlags,
		ProtoFoldFlag.AppLayer
		);

	ProtoRange range = { baseOffset, sizeof(io.DnsHdr), ProtoColor.AppLayer }
	protoRangeArray.append(range);

	io.DnsHdr const* hdr = (io.DnsHdr const*)p;
	p += sizeof(io.DnsHdr);
	size -= sizeof(io.DnsHdr);

	representation.m_lineAttr.m_backColor = ProtoColor.Payload;

	void const* payload = p;
	void const* end = p + size;
	size_t payloadOffset = baseOffset + sizeof(io.DnsHdr);

	if (foldFlags & ProtoFoldFlag.Payload)
	{
		representation.addHyperText($"[\e^-%1\e[34m-\e[m] \e[1mDNS\e[m RRs"(ProtoFoldFlag.Payload));

		std.StringBuilder string;

		if (hdr.m_questionRrCount)
			p = representQuestionRrTableDetails(
				representation,
				string,
				"Question RRs",
				hdr.m_questionRrCount,
				p,
				end,
				hdr,
				baseOffset
				);

		if (hdr.m_answerRrCount)
			p = representAnswerRrTableDetails(
				representation,
				string,
				"Answer RRs",
				hdr.m_answerRrCount,
				p,
				end,
				hdr,
				baseOffset
				);

		if (hdr.m_authorityRrCount)
			p = representAnswerRrTableDetails(
				representation,
				string,
				"Authority RRs",
				hdr.m_authorityRrCount,
				p,
				end,
				hdr,
				baseOffset
				);

		if (hdr.m_additionalRrCount)
			p = representAnswerRrTableDetails(
				representation,
				string,
				"Additional RRs",
				hdr.m_additionalRrCount,
				p,
				end,
				hdr,
				baseOffset
				);
	}
	else
	{
		for (size_t i = 0; p < end && i < hdr.m_questionRrCount; i++)
			p = scanDnsQuestionRr(p, end);

		size_t answerCount =
			hdr.m_answerRrCount +
			hdr.m_authorityRrCount +
			hdr.m_additionalRrCount;

		for (size_t i = 0; p < end && i < answerCount; i++)
			p = scanDnsAnswerRr(p, end);

		representation.addHyperText(
			$"[\e^+%1\e[34m+\e[m] \e[1mDNS\e[m RRs (%2 items, %3 bytes)"(
			ProtoFoldFlag.Payload,
			hdr.m_questionRrCount + answerCount,
			p - payload
			));
	}

	range = { payloadOffset, p - payload, ProtoColor.Payload };
	protoRangeArray.append(range);
}

void const* DnsProtoParser.formatName(
	std.StringBuilder* string,
	void const* p,
	void const* end,
	void const* base
	)
{
	string.clear();
	p = io.formatDnsName(string, p, end, base);
	decodeName(string);
	return p;
}

void const* DnsProtoParser.formatNameLink(
	std.StringBuilder* string,
	void const* p,
	void const* end,
	void const* base,
	size_t baseOffset
	)
{
	void const* p0 = p;
	p = formatName(string, p, end, base);
	size_t offset = baseOffset + p0 - base;
	size_t length = p - p0;
	string.insert(0, $"\e^!0x%x,%2,%3\e[34m"(NetSnifferLogPartCode.RawPacket, offset, length));
	string.append("\e[m");
	return p;
}

char const* DnsProtoParser.getAnswerRrValueString(
	std.StringBuilder* string,
	io.DnsType type,
	void const* p,
	void const* end,
	void const* base
	)
{
	char const* resultString;

	switch (type)
	{
	case io.DnsType.A:
		resultString = ((io.Address_ip4 const*)p).getString();
		break;

	case io.DnsType.AAAA:
		resultString = ((io.Address_ip6 const*)p).getString();
		break;

	case io.DnsType.CNAME:
		formatName(string, p, end, base);
		resultString = string.m_sz;
		break;
	}

	return resultString;
}

void const* scanDnsQuestionRr(
	void const* p,
	void const* end
	)
{
	p = io.scanDnsName(p, end);
	return end - p < sizeof(io.DnsQuestionRrSuffix) ? end : p + sizeof(io.DnsQuestionRrSuffix);
}

void const* scanDnsAnswerRr(
	void const* p,
	void const* end
	)
{
	p = io.scanDnsName(p, end);
	return end - p < sizeof(io.DnsAnswerRrSuffix) ? end : p + sizeof(io.DnsAnswerRrSuffix);
}

void const* DnsProtoParser.representQuestionRrTableDigest(
	log.Representation* representation,
	std.StringBuilder* string,
	size_t count,
	void const* p,
	void const* end,
	void const* base
	)
{
	size_t previewLineIdx = -1;
	if (count > PreviewLineCount)
	{
		count = PreviewLineCount;
		previewLineIdx = PreviewLineCount - 1;
	}

	for (size_t i = 0; p < end && i < count; i++)
	{
		if (i == previewLineIdx)
			representation.m_lineAttr.m_flags |= log.LineAttrFlags.Preview;

		p = representQuestionRrDigest(representation, string, p, end, base);
	}

	return p;
}

void const* DnsProtoParser.representAnswerRrTableDigest(
	log.Representation* representation,
	std.StringBuilder* string,
	size_t questionCount,
	size_t answerCount,
	void const* p,
	void const* end,
	void const* base
	)
{
	size_t endIdx = questionCount + answerCount;
	size_t previewLineIdx = -1;

	if (endIdx > PreviewLineCount)
	{
		endIdx = PreviewLineCount;
		previewLineIdx = PreviewLineCount - 1;
	}

	for (size_t i = questionCount; p < end && i < endIdx; i++)
	{
		if (i == previewLineIdx)
			representation.m_lineAttr.m_flags |= log.LineAttrFlags.Preview;

		p = representAnswerRrDigest(representation, string, p, end, base);
	}

	return p;
}

void const* DnsProtoParser.representQuestionRrTableDetails(
	log.Representation* representation,
	std.StringBuilder* string,
	char const* name,
	size_t count,
	void const* p,
	void const* end,
	void const* base,
	size_t baseOffset
	)
{
	representation.addHyperText($"    $name:");

	for (size_t i = 0; p < end && i < count; i++)
		p = representQuestionRrDetails(representation, string, i, p, end, base, baseOffset);

	return p;
}

void const* DnsProtoParser.representAnswerRrTableDetails(
	log.Representation* representation,
	std.StringBuilder* string,
	char const* name,
	size_t count,
	void const* p,
	void const* end,
	void const* base,
	size_t baseOffset
	)
{
	representation.addHyperText($"    $name:");

	for (size_t i = 0; p < end && i < count; i++)
		p = representAnswerRrDetails(representation, string, i, p, end, base, baseOffset);

	return p;
}

void const* DnsProtoParser.representQuestionRrDigest(
	log.Representation* representation,
	std.StringBuilder* string,
	void const* p,
	void const* end,
	void const* base
	)
{
	if (p >= end)
		return end;

	p = formatName(string, p, end, base);
	if (end - p < sizeof(io.DnsQuestionRrSuffix))
	{
		representation.addHyperText($"Question:\t%1"(string.m_sz));
		return end;
	}

	io.DnsQuestionRrSuffix const* suffix = (io.DnsQuestionRrSuffix const*)p;

	representation.addHyperText($"Question:\t%1 (%2)"(
		string.m_sz,
		typeof(io.DnsType).getValueString(&suffix.m_type)
		));

	return p + sizeof(io.DnsQuestionRrSuffix);
}

void const* DnsProtoParser.representQuestionRrDetails(
	log.Representation* representation,
	std.StringBuilder* string,
	size_t i,
	void const* p,
	void const* end,
	void const* base,
	size_t baseOffset
	)
{
	if (p >= end)
	{
		representation.addHyperText($"%(1;10d): [INCOMPLETE-RR]"(i + 1));
		return end;
	}

	p = formatNameLink(string, p, end, base, baseOffset);

	if (end - p < sizeof(io.DnsQuestionRrSuffix))
	{
		representation.addHyperText(
			$"%(1;10d): Name: %2\n"
			$"%(3;10c)  [INCOMPLETE-RR]"(
			i + 1,
			string.m_sz,
			' '
			));

		return end;
	}

	representation.addHyperText($"%(1;10d): Name:\t%2"(i + 1, string.m_sz));

	log.representStruct(
		representation,
		typeof(io.DnsQuestionRrSuffix),,
		NetSnifferLogPartCode.RawPacket,
		p,
		baseOffset + p - base,
		0,
		0,
		2
		);

	return p + sizeof(io.DnsQuestionRrSuffix);
}

void const* DnsProtoParser.representAnswerRrDigest(
	log.Representation* representation,
	std.StringBuilder* string,
	void const* p,
	void const* end,
	void const* base
	)
{
	if (p >= end)
		return end;

	p = formatName(string, p, end, base);

	if (end - p < sizeof(io.DnsAnswerRrSuffix))
	{
		representation.addHyperText($"Answer:\t%1"(string.m_sz));
		return end;
	}

	io.DnsAnswerRrSuffix const* suffix = (io.DnsAnswerRrSuffix const*)p;
	p += sizeof(io.DnsAnswerRrSuffix);

	if (end - p < suffix.m_dataLength)
	{
		representation.addHyperText($"Answer:\t%1 (%2)"(
			string.m_sz,
			typeof(io.DnsType).getValueString(&suffix.m_type)
			));

		return end;
	}

	char const* valueString = getAnswerRrValueString(string, suffix.m_type, p, end, base);

	representation.addHyperText($"Answer:\t%1 (%2%3%4)"(
		string.m_sz,
		typeof(io.DnsType).getValueString(&suffix.m_type),
		valueString ? ": " : null,
		valueString
		));

	return p + suffix.m_dataLength;
}

void const* DnsProtoParser.representAnswerRrDetails(
	log.Representation* representation,
	std.StringBuilder* string,
	size_t i,
	void const* p,
	void const* end,
	void const* base,
	size_t baseOffset
	)
{
	if (p >= end)
	{
		representation.addHyperText($"%(1;10d): [INCOMPLETE-RR]"(i + 1));
		return end;
	}

	p = formatNameLink(string, p, end, base, baseOffset);

	if (end - p < sizeof(io.DnsAnswerRrSuffix))
	{
		representation.addHyperText(
			$"%(1;10d): Name: %2\n"
			$"%(3;10c)  [INCOMPLETE-RR]"(
			i + 1,
			string.m_sz,
			' '
			));

		return end;
	}

	representation.addHyperText($"%(1;10d): Name:\t%2"(i + 1, string.m_sz));

	log.representStruct(
		representation,
		typeof(io.DnsAnswerRrSuffix),,
		NetSnifferLogPartCode.RawPacket,
		p,
		baseOffset + p - base,
		0,
		0,
		2
		);

	io.DnsAnswerRrSuffix const* suffix = (io.DnsAnswerRrSuffix const*)p;
	p += sizeof(io.DnsAnswerRrSuffix);

	if (end - p < suffix.m_dataLength)
	{
		representation.addHyperText($"%(1;12c)[INCOMPLETE-RR]"(' '));
		return end;
	}

	char const* valueString = getAnswerRrValueString(string, suffix.m_type, p, end, base);

	if (valueString)
		representation.addHyperText(
			$"%(1;12c)Value:\t\e^!0x%x,%3,%4\e[34m%5"(
			' ',
			NetSnifferLogPartCode.RawPacket,
			baseOffset + p - base,
			suffix.m_dataLength,
			valueString
			));

	return p + suffix.m_dataLength;
}

//..............................................................................

enum NetBiosNameType: char
{
	Workstation               = 0x00,
	MessengerService          = 0x03,
	RemoteAccessService       = 0x06,
	DomainMasterBrowser       = 0x1b,
	DomainController          = 0x1c,
	MasterBrowser             = 0x1d,
	FileService               = 0x20,
	RemoteAccessServiceClient = 0x21,
}

class NbnsProtoParser: DnsProtoParser
{
	construct()
	{
		basetype.construct("NBNS");
	}

protected:
	override decodeName(std.StringBuilder* string)
	{
		size_t originalLength = string.m_length;
		size_t length = originalLength / 2;
		string.append(' ',  length);

		for (size_t i = 0, j = 0; i < length; i++, j += 2)
			string[originalLength + i] = ((string[j] - 'A') << 4) | (string[j + 1] - 'A');

		string.remove(0, originalLength);

		if (length == 16) // last character is the type
		{
			char type = string[15];
			string.chop(1);
			string.trimRight();
			string.append($" (%1)"(typeof(NetBiosNameType).getValueString(type, "0x%02x")));
		}
	}
}

//..............................................................................

construct()
{
	static DnsProtoParser dns;
	static NbnsProtoParser nbns;
	static DnsProtoParser mdns("MDNS");
	static DnsProtoParser llmnr("LLMNR");

	g_udpPortAppProtoParserMap[53] = dns;
	g_udpPortAppProtoParserMap[137] = nbns;
	g_udpPortAppProtoParserMap[5353] = mdns;
	g_udpPortAppProtoParserMap[5355] = llmnr;
}

//..............................................................................
