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

#include <arpa/inet.h>

#include "dns.h"

/* DNS packet:
  +--------+-------+-----------------------------------------------+
  | Offset | Bytes  | Field                                        |
  +--------+--------+----------------------------------------------+
  | 0      | 2      | Transaction ID                               |
  | 2      | 2      | Flags                                        |
  | 4      | 2      | QDCOUNT (Number of questions)                |
  | 6      | 2      | ANCOUNT (Number of answers)                  |
  | 8      | 2      | NSCOUNT (Number of authority RRs)            |
  | 10     | 2      | ARCOUNT (Number of additional RRs)           |
  +--------+--------+----------------------------------------------+
  | 12     | var    | Questions (name + qtype + qclass)            |
  +--------+--------+----------------------------------------------+
  | ??     | var    | Answers (name+type+class+TTL+RDLENGTH+RDATA) |
  +--------+--------+----------------------------------------------+
  | ??     | var    | Authority RRs (same format)                  |
  +--------+--------+----------------------------------------------+
  | ??     | var    | Additional RRs (same format)                 |
  +--------+--------+----------------------------------------------+
*/
size_t parse_dns_udp(const uint8_t *payload,
					 size_t len,
					 struct dns_question *out, size_t max_qs) {
	if (len < 12) {
		// not long enough
		return 0;
	}

	uint16_t qdcount = ntohs(*(uint16_t *)(payload + 4));
	if (qdcount == 0) {
		return 0;
	}
	if (qdcount > max_qs) {
		qdcount = max_qs; // truncate safely
	}

	const uint8_t *ptr = payload + 12;
	const uint8_t *end = payload + len;

	size_t parsed = 0;
	for (int q = 0; q < qdcount && ptr < end; q++) {
		char name[MAX_DNS_NAME_LEN];
		int name_len = 0;

		while (ptr < end && *ptr != 0 && name_len < MAX_DNS_NAME_LEN - 1) {
			uint8_t label_len = *ptr++;

			if (label_len == 0 || ptr + label_len > end) {
				return 0;
			}

			if (name_len > 0) {
				name[name_len++] = '.';
			}

			memcpy(name + name_len, ptr, label_len);
			name_len += label_len;
			ptr += label_len;
		}

		name[name_len] = '\0';
		ptr++; // skip null terminator

		if (ptr + 4 > end) {
			return 0;
		}

		uint16_t qtype = ntohs(*(uint16_t *)(ptr));
		ptr += 2;
		ptr += 2; // qclass

		strncpy(out[parsed].name, name, MAX_DNS_NAME_LEN);
		out[parsed].name[MAX_DNS_NAME_LEN - 1] = '\0';
		out[parsed].qtype = qtype;

		parsed++;
	}

	return parsed;
}

// resolve DNS query type to human-readable string.
// TODO static buffer may yield incorrect resultsif multiple unknown
//      DNS types are encountered in a single message.
const char *dns_type_to_string(uint16_t type) {
	switch(type) {
	case 1:   return "A";
	case 2:   return "NS";
	case 5:   return "CNAME";
	case 6:   return "SOA";
	case 12:  return "PTR";
	case 15:  return "MX";
	case 16:  return "TXT";
	case 28:  return "AAAA";
	case 33:  return "SRV";
	case 64:  return "SVCB"; // RFC 9460 https://www.rfc-editor.org/rfc/rfc9460
	case 65:  return "HTTPS"; // RFC 9460 https://www.rfc-editor.org/rfc/rfc9460
	case 255: return "ANY";
	default: // return TYPEXXXX if type isn't known
		static char buf[16];
		snprintf(buf, sizeof(buf), "TYPE%d", type);
		return buf;
	}
}