/* Source code for dnsmeddler
 * This is a little DNS pass-through server
 * It allows the user to creatively change DNS queries
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include "scolor.h"

// ChatGPT wrote this
struct dns_header {
	unsigned short id;
	unsigned char qr : 1;
	unsigned char opcode : 4;
	unsigned char aa : 1;
	unsigned char tc : 1;
	unsigned char rd : 1;
	unsigned char ra : 1;
	unsigned char z : 3;
	unsigned char rcode : 4;
	unsigned short qdcount; // ChatGPT thought these didn't need to be converted
	unsigned short ancount;
	unsigned short nscount;
	unsigned short arcount;
} __attribute__((packed));

// ChatGPT wrote these two functions for me
void dns_header_to_network(struct dns_header *header) {
	header->id = htons(header->id);
	header->qdcount = htons(header->qdcount);
	header->ancount = htons(header->ancount);
	header->nscount = htons(header->nscount);
	header->arcount = htons(header->arcount);
}

void dns_header_from_network(struct dns_header *header) {
	header->id = ntohs(header->id);
	header->qdcount = ntohs(header->qdcount);
	header->ancount = ntohs(header->ancount);
	header->nscount = ntohs(header->nscount);
	header->arcount = ntohs(header->arcount);
}

// ChatGPT wrote this too
void print_dns_header(const struct dns_header *header) {
	printf("ID: %u\n", header->id);
	printf("QR: %u\n", header->qr);
	printf("Opcode: %u\n", header->opcode);
	printf("AA: %u\n", header->aa);
	printf("TC: %u\n", header->tc);
	printf("RD: %u\n", header->rd);
	printf("RA: %u\n", header->ra);
	printf("Z: %u\n", header->z);
	printf("RCODE: %u\n", header->rcode);
	printf("QDCOUNT: %u\n", header->qdcount);
	printf("ANCOUNT: %u\n", header->ancount);
	printf("NSCOUNT: %u\n", header->nscount);
	printf("ARCOUNT: %u\n", header->arcount);
}

// Again, ChatGPT wrote this function
/* But it didn't work!  Parsed like this:
 * query for www.microsoft.com came out as www.icrosoftcom (113 bytes)
 * query for www.google.com came out as www.ooglecom (107 bytes)
*/
/*
size_t parse_domain_name(const char *query_ptr, size_t query_len, char *out_buf, size_t out_buf_len) {
	size_t i = 0, j = 0;
	unsigned int label_len = 0;
	int compressed = 0;

	// Parse the domain name labels
	while (i < query_len) {
		// Check for compression
		if ((query_ptr[i] & 0xC0) == 0xC0) {
			// Compressed pointer
			if (i + 1 >= query_len) {
				// Malformed packet
				return 0;
			}
			unsigned int offset = ((query_ptr[i] & 0x3F) << 8) | query_ptr[i + 1];
			if (offset >= query_len) {
				// Malformed packet
				return 0;
			}
			if (!compressed) {
				compressed = 1;
				label_len += 2; // Add 2 bytes for the compressed pointer
			}
			query_ptr = query_ptr + offset - 0xC000;
			query_len = query_len - offset + 0xC000;
			i = 0;
		} else {
			// Regular label
			label_len = query_ptr[i];
			i++;
			if (label_len == 0) {
				break;
			}
			if (!compressed) {
				label_len += 1; // Add 1 byte for the label length field
			}
			if (label_len > out_buf_len - j - 1) {
				// Output buffer too small
				return 0;
			}
			memcpy(out_buf + j, query_ptr + i, label_len);
			i += label_len;
			j += label_len;
			if (!compressed) {
				out_buf[j - 1] = '.'; // Replace label length with dot separator
			}
		}
	}

	if (!compressed && j > 0) {
		// Remove trailing dot separator
		out_buf[j - 1] = '\0';
		j--;
	}

	return j;
}
*/

/* ChatGPT added a fix, in this function */
/* We're closer!
 * www.google.com -> ww.googl.co
 * www.microsoft.com -> ww.microsof.co
 * Missing the last letter of each bit of text, basically
 */
/*
size_t parse_domain_name(const char *query_ptr, size_t query_len, char *out_buf, size_t out_buf_len) {
	size_t i = 0, j = 0;
	unsigned int label_len = 0;
	int compressed = 0;

	// Parse the domain name labels
	while (i < query_len) {
		// Check for compression
		if ((query_ptr[i] & 0xC0) == 0xC0) {
			// Compressed pointer
			if (i + 1 >= query_len) {
				// Malformed packet
				return 0;
			}
			unsigned int offset = ((query_ptr[i] & 0x3F) << 8) | query_ptr[i + 1];
			if (offset >= query_len) {
				// Malformed packet
				return 0;
			}
			if (!compressed) {
				compressed = 1;
				label_len += 2; // Add 2 bytes for the compressed pointer
			}
			query_ptr = query_ptr + offset - 0xC000;
			query_len = query_len - offset + 0xC000;
			i = 0;
		} else {
			// Regular label
			label_len = query_ptr[i];
			i++;
			if (label_len == 0) {
				break;
			}
			if (!compressed) {
				label_len += 1; // Add 1 byte for the label length field
			}
			if (label_len > out_buf_len - j - 1) {
				// Output buffer too small
				return 0;
			}
			memcpy(out_buf + j, query_ptr + i, label_len - 1);
			i += label_len - 1;
			j += label_len - 1;
			if (!compressed) {
				out_buf[j - 1] = '.'; // Replace label length with dot separator
			}
		}
	}

	if (!compressed && j > 0) {
		// Remove trailing dot separator
		out_buf[j - 1] = '\0';
		j--;
	}

	return j;
}
*/

/* Here's ChatGPT's third try! */
/* As far as I can tell, this works correctly
 * To be perfectly honest, I didn't read it
 */
size_t parse_domain_name(const char *query_ptr, size_t query_len, char *out_buf, size_t out_buf_len) {
	size_t i = 0, j = 0;
	unsigned int label_len = 0;
	int compressed = 0;

	// Parse the domain name labels
	while (i < query_len) {
		// Check for compression
		if ((query_ptr[i] & 0xC0) == 0xC0) {
			// Compressed pointer
			if (i + 1 >= query_len) {
				// Malformed packet
				return 0;
			}
			unsigned int offset = ((query_ptr[i] & 0x3F) << 8) | query_ptr[i + 1];
			if (offset >= query_len) {
				// Malformed packet
				return 0;
			}
			if (!compressed) {
				compressed = 1;
				label_len += 2; // Add 2 bytes for the compressed pointer
			}
			query_ptr = query_ptr + offset - 0xC000;
			query_len = query_len - offset + 0xC000;
			i = 0;
		} else {
			// Regular label
			label_len = query_ptr[i];
			i++;
			if (label_len == 0) {
				break;
			}
			if (!compressed) {
				label_len += 1; // Add 1 byte for the label length field
			}
			if (label_len > out_buf_len - j - 1) {
				// Output buffer too small
				return 0;
			}
			memcpy(out_buf + j, query_ptr + i, label_len - 1);
			i += label_len - 1;
			j += label_len - 1;
			if (!compressed && i < query_len) {
				// Add dot separator between labels
				out_buf[j] = '.';
				j++;
			}
		}
	}

	if (!compressed && j > 0) {
		// Remove trailing dot separator
		out_buf[j - 1] = '\0';
		j--;
	}

	return j;
}

/* ChatGPT wrote all of this except for the verbose system */
#define BUF_SIZE 2048 // Probably bigger than the MTU
char verbose = 0;
int main() {
	// Create a UDP socket
	int sockfd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
	if (sockfd == -1) {
		perror("socket");
		exit(EXIT_FAILURE);
	}

	// Bind the socket to port 53
	struct sockaddr_in addr = {0};
	addr.sin_family = AF_INET;
	addr.sin_port = htons(53);
	addr.sin_addr.s_addr = INADDR_ANY;
	if (bind(sockfd, (struct sockaddr *) &addr, sizeof(addr)) == -1) {
		perror("bind");
		exit(EXIT_FAILURE);
	}
	/*
	 * Receive a DNS query
	 * Pass the DNS query to a real DNS server 
	 * Receive the result from the real DNS server
	 * Return the result to the client
	 */
	for(;;){
main_loop_top:
		printf(DGREEN("---------------- New Iteration ---------------------\n"));
		// Wait for a DNS message to arrive
		char buf[BUF_SIZE];
		struct sockaddr_in sender_addr = {0};
		socklen_t sender_len = sizeof(sender_addr);
		ssize_t recv_len = recvfrom(sockfd, buf, sizeof(buf), 0, (struct sockaddr *) &sender_addr, &sender_len);
		if (recv_len == -1) {
			perror("recvfrom");
			exit(EXIT_FAILURE);
		}
	
		// Parse the DNS header
		struct dns_header header;
		memcpy(&header, buf, sizeof(header));
		dns_header_from_network(&header);
	
		// Print the DNS header fields
		printf("Received DNS message from %s:%u\n", inet_ntoa(sender_addr.sin_addr), ntohs(sender_addr.sin_port));
		if(verbose) print_dns_header(&header);
	
	 // Parse the question section
reparse:;
		char *query_ptr = buf + sizeof(header);
		size_t remaining_len = recv_len - sizeof(header);
		for (unsigned int i = 0; i < header.qdcount; i++) {
			// Parse the domain name
			char domain_name[256];
			size_t domain_name_len = parse_domain_name(query_ptr, remaining_len, domain_name, sizeof(domain_name));
			if (domain_name_len == 0) {
				fprintf(stderr, "Error: Failed to parse domain name in question %u\n", i);
				exit(EXIT_FAILURE);
			}
	
			// Print the domain name and size
			printf(PURPLE("Question %u: ")CYAN("%s")" (%lu bytes)\n", i, domain_name, domain_name_len);
			if(strstr(domain_name, "google")){
				printf(YELLOW("Warning:  Query is for Google\n"));
				//goto main_loop_top;
				memcpy(query_ptr + 5, "reddit", 6);
				printf(YELLOW("Changed query to reddit\n"));
				goto reparse;
			}
	
			// Advance the query pointer and update the remaining length
			query_ptr += domain_name_len + 4; // Add 4 bytes for the query type and query class fields
			remaining_len -= domain_name_len + 4;
		}
		/* At this point, ChatGPT got really slow and stopped, so I guess it went down
		 * So from here on I wrote all of it except the close at the end 
		 * You can tell my code from ChatGPT's code because my code is messy and uncommented
		 */
		struct sockaddr_in realdns_addr = {AF_INET, htons(53), {0x01010101}}; // 0x01010101 = 1.1.1.1
		int realdns_fd = socket(AF_INET, SOCK_DGRAM, 0);
		if(verbose) printf("Sending off the query to %s port %u\n", inet_ntoa(realdns_addr.sin_addr), ntohs(realdns_addr.sin_port));
		ssize_t sendlen = sendto(realdns_fd, buf, recv_len, 0, (struct sockaddr*)(&realdns_addr), sizeof(realdns_addr));
		if(verbose) printf("Query send length was %ld bytes\n", sendlen);
		if(sendlen < 1)
			printf(RED("Something went wrong with sending\n"));
		char response[2048];
		socklen_t ara_len = sizeof(realdns_addr);
		if(verbose) printf("Waiting for a response...\n");
		// If we don't hear back in 1 second, try again!
		ssize_t response_length = recvfrom(realdns_fd, response, 2048, 0, (struct sockaddr*)&realdns_addr, &ara_len);

		if(verbose) printf("Response received, length was %ld bytes\n", response_length);
		struct dns_header response_header;
		memcpy(&response_header, response, sizeof(response_header));
		dns_header_from_network(&response_header);
		if(verbose) print_dns_header(&response_header);
		if(verbose) printf("Sending reply to original client\n");
		ssize_t replylen = sendto(sockfd, response, response_length, 0, (struct sockaddr*)(&sender_addr), sender_len);
		if(verbose) printf("Response sent, length was %lu bytes\n", replylen);
		if(replylen < 1)
			printf(RED("Something went wrong with our final reply\n"));
	
	}
	// Clean up
	close(sockfd);
	return 0;
}