#dns.c

/*
 * author: kay ~ irc.nullnetwork.net
 * date: 31/08/2006
 *
 * simple DNS query implementation. it's not completed to what i want it to
 * do once i've completely finished it, but it is working properly now. mail
 * me if you spot anything error wise. i'm already aware of certain
 * insecurities in the code. this is just to show how it works at the moment.n	
 */


#include <sys/types.h>
#include <sys/socket.h>
#include <stdio.h>
#include <netinet/in.h>
#include <string.h>
#include <stdlib.h>

#define PORT	53
#define NS	"195.92.195.95"
#define BUFLEN	511

#define QTYPE_A 1
#define QTYPE_NS 2
#define QTYPE_MD 3
#define QTYPE_MF 4
#define QTYPE_CNAME 5
#define QTYPE_SOA 6
#define QTYPE_MB 7
#define QTYPE_MG 8
#define QTYPE_MR 9
#define QTYPE_NULL 10
#define QTYPE_WKS 11
#define QTYPE_PTR 12
#define QTYPE_HINFO 13
#define QTYPE_MINFO 14
#define QTYPE_MX 15
#define QTYPE_TXT 16
#define QTYPE_AXFR 252
#define QTYPE_MAILB 253
#define QTYPE_MAILA 254
#define QTYPE_OTHER 255
#define QCLASS_IN 1
#define QCLASS_CS 2
#define QCLASS_CH 3
#define QCLASS_HS 4
#define QCLASS_OTHER 255

#define PGREEN(x) fwrite("\x1B[01;32m" x "\x1B[00m\n", 1, 14+sizeof(x), stdout);
#define PRED(x) fwrite("\x1B[01;31m" x "\x1B[00m\n", 1, 14+sizeof(x), stdout);
#define PGREY(x) fwrite("\x1B[01;30m" x "\x1B[00m\n", 1, 14+sizeof(x), stdout);
#define PWHITE(x) fwrite("\x1B[01;29m" x "\x1B[00m\n", 1, 14+sizeof(x), stdout);
#define PYELLOW(x) fwrite("\x1B[01;33m" x "\x1B[00m\n", 1, 14+sizeof(x), stdout);
#define PBLUE(x) fwrite("\x1B[01;34m" x "\x1B[00m\n", 1, 14+sizeof(x), stdout);
#define PPURPLE(x) fwrite("\x1B[01;35m" x "\x1B[00m\n", 1, 14+sizeof(x), stdout);

struct dns_header {
	short int id; // 2 bytes
	short int flags;
	/* bit n | meaning
	   0     | QR: query responce
	   1-4   | Opcode: query opcode (0 - standard query, 1 - inverse query, 2 = server status req., 3-15 reserved)
	   5     | AA: authoritive answer
	   6     | TC: Truncation
	   7     | RD: Recursion desired
	   8     | RA: recursion available
	   9-11  | Z: Reserved
	   12-15 | RCODE: responce code (0 - no error, 1 - format error, 2 - server failure, 3 - name error, 4 - not implemented, 5 - refused, 6-16 - reserved)
	*/
	short int qdcount;
	short int ancount;
	short int nscount;
	short int arcount;
};

struct dns_question {
	char * qname; // query name
	short int qtype;
	short int qclass;
};

struct dns_resource {

	char * name; // if name & 0xC0 == 0xC0 then its a pointer to name & 0x3FFF
	short int type;
	short int class;
	unsigned int ttl;
	unsigned short int rdlength;
	char * rdata; // note this is variable length, given by rdlength
};


typedef struct dns_header dns_header;
typedef struct dns_question dns_question;
typedef struct dns_resource dns_answer;
typedef struct dns_resource dns_authority;
typedef struct dns_resource dns_additional;

#define SET_QR(x) x|=0x8
#define UNSET_QR(x) x~=0x8

struct rdata_mx {
	short int preference;
	char * exchange;
};

struct rdata_soa {
	char * mname;
	char * rname;
	unsigned int serial;
	int refresh;
	int retry;
	int expire;
	unsigned int minimum;
};


int sendallto(int fd, char * data, int datalen, struct sockaddr_in * dest) {
int sendlen = datalen, num_of_bytes;	
	while (sendlen > 0) {

		if ((num_of_bytes = sendto(fd, data + (datalen - sendlen), sendlen, 0, (struct sockaddr *)dest, sizeof(struct sockaddr))) == -1) {
			return -1;
		}
		sendlen -= num_of_bytes;
	}
	return 0;
}


int make_query(int fd, struct sockaddr_in dest_addr, dns_question query) {
dns_header qhead;
char buf[BUFLEN];
int buflen = 0;
int octetlen = 0;
char * dotptr = NULL;
char * hostptr = query.qname;
		
	memset(buf, 0, BUFLEN);
	
	qhead.id = htons(0x1337);
	qhead.flags = htons(0x0100);
	qhead.qdcount = htons(0x0001);
	qhead.ancount = 0;
	qhead.nscount = 0;
	qhead.arcount = 0;
	
	buflen = sizeof(qhead);
	memcpy(buf, (void*)&qhead, buflen);
	
	do {
		dotptr = strchr(hostptr, '.');
		
		if (!dotptr)
			octetlen = strchr(hostptr,'\0') - hostptr;
		else 
			octetlen = dotptr - hostptr;
		
		buf[buflen++] = octetlen;
		strncpy(buf+buflen, hostptr, octetlen);
		buflen += octetlen;
		hostptr = dotptr + 1;
	} while(dotptr);
	
	memset(buf+buflen, 0, 2);
	buflen+=2;
	//qtype
	query.qtype = query.qtype;
	query.qclass = query.qclass;
	memcpy((void*)buf+buflen, (void*)&(query.qtype), sizeof(short));
	buflen += 2;
	//qclass
	memcpy((void*)buf+buflen, (void*)&(query.qclass), sizeof(short));
	buflen += 2;
	
	sendallto(fd, buf, buflen, &dest_addr);
	

}

char * translate_qtype(int qt) {
	switch(qt) {
	case QTYPE_A: return "A";
	case QTYPE_NS: return "NS";
	case QTYPE_MD: return "MD";
	case QTYPE_MF: return "MF";
	case QTYPE_CNAME: return "CNAME";
	case QTYPE_SOA: return "SOA";
	case QTYPE_MB: return "MB";
	case QTYPE_MG: return "MG";
	case QTYPE_MR: return "MR";
	case QTYPE_NULL: return "NULL";
	case QTYPE_WKS: return "WKS";
	case QTYPE_PTR: return "PTR";
	case QTYPE_HINFO: return "HINFO";
	case QTYPE_MINFO: return "MINFO";
	case QTYPE_MX: return "MX";
	case QTYPE_TXT: return "TXT";
	case QTYPE_AXFR: return "AXFR";
	case QTYPE_MAILB: return "MAILB";
	case QTYPE_MAILA: return "MAILA";
	case QTYPE_OTHER: return "*";
	default: return NULL;
	}
}

char * translate_qclass(int qc) {
	switch(qc) {
	case QCLASS_IN: return "IN";
	case QCLASS_CS: return "CS";
	case QCLASS_CH: return "CH";
	case QCLASS_HS: return "HS";
	case QCLASS_OTHER: return "*";
	default: return NULL;
	}
	
}

char * translate_rcode(int rc) {
	switch(rc) {
	case 0: return "No Error";
	case 1: return "Format Error";
	case 2: return "Server Failure";
	case 3: return "Name Error";
	case 4: return "Not Implemented";
	case 5: return "Refused";
	default:
		return "Reserved";
	}
}

char * translate_opcode(int oc) {
	switch(oc) {
	case 0: return "Standard Query";
	case 1: return "Inverse Query";
	case 2: return "Server Status Request";
	default:
		return "Reserved";
	}
}

// ---[data extraction rtouines]---

char * get_charstring(char * buf, int * offset) {
char * str = (char*)malloc(256);
int length;
	
	length = buf[(*offset)++];
	strncpy(str, buf+(*offset), length);
	*offset += length;
	return str;
}

char * get_name(char * data, int * start) {
char * name = (char*)malloc(64);
int name_len = 0;
short int offset = *start;
unsigned short int octetlen = 0;
struct qname_a * qname;
short int ptroffset = 0;
int set_start = 1;
	if(!name) return NULL;
	memset(name, 0, 64);
	do {
		
		if ((data[offset] & 0xC0) == 0xC0) {
			ptroffset = 0;
			ptroffset |= (data[offset++] & 0x3F);
			ptroffset <<= 8;
			ptroffset |= (data[offset++] & 0xFF);
			if (set_start) *start = offset;
			offset = ptroffset;
			set_start = 0;
		}
		octetlen = data[offset++];
		if (octetlen) {
			if (name_len) name[name_len++] = '.';
			strncpy(name + name_len, data + offset, octetlen);
			name_len += octetlen;
			offset += octetlen;
		}
	} while(octetlen);
	if (set_start) *start = offset;
	return name;
}

// polymorphism, where for art thou

short int readshortintvalue(char * buf, int * offset) {
short int bufint = 0;
	memcpy((void*)&(bufint), buf + *offset, sizeof(short int));
	*offset += sizeof(short int);
	return ntohs(bufint);
}
int readintvalue(char * buf, int * offset) {
int bufint = 0;
	memcpy((void*)&(bufint), buf + *offset, sizeof(int));
	*offset += sizeof(int);
	return ntohl(bufint);
}

// ---[ rdata functions ]---

struct rdata_mx * get_rdata_mx(char * buf, int * offset) {
struct rdata_mx * mx = (struct rdata_mx *)calloc(1, sizeof(struct rdata_mx));
	if (!mx) return NULL;
	memcpy((void*)&(mx->preference), buf + *offset, sizeof(short));
	*offset += sizeof(short);
	mx->exchange = get_name(buf, offset);
	return mx;
}


struct rdata_soa * get_rdata_soa(char * buf, int * offset) {
struct rdata_soa * soa = (struct rdata_soa *)calloc(1, sizeof(struct rdata_soa));
	if (!soa) return NULL;
	soa->mname = get_name(buf, offset);
	soa->rname = get_name(buf, offset);
	soa->serial = readintvalue(buf, offset);
	soa->refresh = readintvalue(buf, offset);
	soa->retry = readintvalue(buf, offset);
	soa->expire = readintvalue(buf, offset);
	soa->minimum = readintvalue(buf, offset);
	
	return soa;
}


struct dns_resource * get_resource(char * buf, int * start) {
struct dns_resource * dns_res = (struct dns_resource *)calloc(1, sizeof(struct dns_resource));
int offset = * start;
char * name, * txt;
int inetaddr;
struct rdata_mx * mx;
struct rdata_soa * soa;

	if(!dns_res) return NULL;
	name = get_name(buf, &offset);
	dns_res->name = name;
	dns_res->type = readshortintvalue(buf, &offset);
	dns_res->class = readshortintvalue(buf, &offset);
	dns_res->ttl = readintvalue(buf, &offset);
	dns_res->rdlength = readshortintvalue(buf, &offset);
	
	printf("dns_res {\n");
	printf("  name='%s'\n", dns_res->name);
	printf("  type=%d (Type: %s)\n", dns_res->type, translate_qtype(dns_res->type));
	printf("  class=%X (Class: %s)\n", dns_res->class, translate_qclass(dns_res->class));
	printf("  ttl=%d seconds\n", dns_res->ttl);
	printf("  rdlength=%X\n", dns_res->rdlength);
	//dns_ans.rdata 
	switch (dns_res->type) {
	case QTYPE_MB:
	case QTYPE_MD:
	case QTYPE_MF:
	case QTYPE_MG:
	case QTYPE_MR:
	case QTYPE_NS:
	case QTYPE_PTR:
	case QTYPE_CNAME:
		// domain-name
		name = get_name(buf, &offset);
		printf("  rdatas='%s' when type is %s\n", name, translate_qtype(dns_res->type));
		dns_res->rdata = name;
		break;
	case QTYPE_WKS:
		//
		break;
	case QTYPE_A:
		memcpy((void*)&(inetaddr), buf + offset, sizeof(int));
		offset += sizeof(int);
		printf("  rdata=0x%X (%s)\n", inetaddr, inet_ntoa(inetaddr));
		dns_res->rdata = (char*)inetaddr;
		break;
	case QTYPE_TXT:
		txt = get_charstring(buf, &offset);
		printf("  rdata='%s'\n", txt);
		break;
	case QTYPE_SOA:
		soa = get_rdata_soa(buf, &offset);
		dns_res->rdata = (char*)soa;
		printf("  rdata {\n");
		printf("    mname='%s'\n", soa->mname);
		printf("    rname='%s'\n", soa->rname);
		printf("    serial=%d\n", soa->serial);
		printf("    refresh=%d seconds\n", soa->refresh);
		printf("    retry=%d seconds\n", soa->retry);
		printf("    expire=%d seconds\n", soa->expire);
		printf("    minimum=%d seconds\n", soa->minimum);
		printf("  }\n");
		break;
	case QTYPE_NULL:
		// this is an experimental type, so I'll just ignore it
		printf("  # rdata ignored due to type being experimental\n");
		break;
	case QTYPE_MX:
		// preference
		// exchange (domain-name)
		mx = get_rdata_mx(buf, &offset);
		dns_res->rdata = (char*)mx;
		printf("  rdata {\n");
		printf("    preference=%d\n", mx->preference);
		printf("    exchange='%s'\n", mx->exchange);
		printf("  }\n");
		break;
	case QTYPE_MINFO:
		// rmailbx (domain-name)
		// emailbx (domain-name)
		break;
	case QTYPE_HINFO:
		// cpu (char-string)
		// os (char-string)
		break;
	default:
		printf("Unknown type\n");
	}
	printf("}\n");
	*start = offset;

}

int read_responce(int fd, struct sockaddr_in dest_addr) {
int num_of_bytes = 0;
char buf[BUFLEN];
unsigned int buflen = 0;
dns_header dns_head;
dns_question dns_quest;
dns_authority * dns_auth;
dns_answer * dns_ans;
dns_additional * dns_add;
int offset = 0;
char * name;
int i;
unsigned short int qr, opcode, aa, tc, rd, ra, z, rcode;


	if ((num_of_bytes = recvfrom(fd, buf, BUFLEN, 0, (struct sockaddr *)&dest_addr, &buflen)) == -1) {
		perror("recvfrom() error");
		return 1;
	}
	printf("%d bytes received\n", num_of_bytes);
	
	memcpy((void*)&dns_head, buf, sizeof(dns_header));
	dns_head.id = ntohs(dns_head.id);
	dns_head.flags = ntohs(dns_head.flags);
	dns_head.qdcount = ntohs(dns_head.qdcount);
	dns_head.ancount = ntohs(dns_head.ancount);
	dns_head.nscount = ntohs(dns_head.nscount);
	dns_head.arcount = ntohs(dns_head.arcount);
	
	qr = dns_head.flags & 0x8000;
	opcode = (dns_head.flags & 0x7800) >> 11;
	aa = dns_head.flags & 0x0400;
	tc = dns_head.flags & 0x0200;
	rd = dns_head.flags & 0x0100;
	ra = dns_head.flags & 0x0080;
	z = (dns_head.flags & 0x0070) >> 4;
	rcode = (dns_head.flags & 0x000F);
	
	printf("dns_head {\n");
	printf("  id=0x%X\n", dns_head.id);
	printf("  flags=0x%hX {\n", dns_head.flags);
	printf("    qr=%u (%s)\n", (qr ? 1 : 0), (qr ? "is a Query Responce" : "is not a Query Responce"));
	printf("    opcode=0x%X (%s)\n", opcode, translate_opcode(opcode));
	printf("    aa=%u (%s)\n", (aa ? 1 : 0), (aa ? "is the Authoritive Server" : "is not the Authoritive Server"));
	printf("    tc=%u%s\n", (tc ? 1 : 0), (tc ? " ( is Truncated)" : ""));
	printf("    rd=%u%s\n", (rd ? 1 : 0), (rd ? " (Recursion Desired)" : ""));
	printf("    ra=%u%s\n", (ra ? 1 : 0), (ra ? " (Recursion available)" : ""));
	printf("    z=0x%X\n", z);
	printf("    rcode=0x%X (%s)\n", rcode, translate_rcode(rcode));
	printf("  }\n");
	printf("  qdcount=%u\n", dns_head.qdcount);
	printf("  ancount=%u\n", dns_head.ancount);
	printf("  nscount=%u\n", dns_head.nscount);
	printf("  arcount=%u\n", dns_head.arcount);
	printf("}\n");
	/* bit n | meaning
	   0     | QR: query responce
	   1-4   | Opcode: query opcode (0 - standard query, 1 - inverse query, 2 = server status req., 3-15 reserved)
	   5     | AA: authoritive answer
	   6     | TC: Truncation
	   7     | RD: Recursion desired
	   8     | RA: recursion available
	   9-11  | Z: Reserved
	   12-15 | RCODE: responce code (0 - no error, 1 - format error, 2 - server failure, 3 - name error, 4 - not implemented, 5 - refused, 6-16 - reserved)
	*/	
	
	// get question information
	offset = sizeof(dns_header);
	
	// lets get qname
	name = get_name(buf, &offset);
	dns_quest.qname = name;
	//add_qname(name, offset);
	
	// get qtype and qclass
	memcpy((void*)&(dns_quest.qtype), buf + offset, sizeof(short int));
	offset += sizeof(short int);
	dns_quest.qtype = ntohs(dns_quest.qtype);
	
int inetaddr;
struct rdata_mx * mx;
	memcpy((void*)&(dns_quest.qclass), buf + offset, sizeof(short int));
	offset += sizeof(short int);
	dns_quest.qclass = ntohs(dns_quest.qclass);
	
	printf("dns_quest {\n");
	printf("  qname='%s'\n", dns_quest.qname);
	printf("  qtype=%d (Type: %s)\n", dns_quest.qtype, translate_qtype(dns_quest.qtype));
	printf("  qclass=%X (Class: %s)\n", dns_quest.qclass, translate_qclass(dns_quest.qclass));
	printf("}\n");
	PBLUE("Answers:");
	for(i =0; i < dns_head.ancount; i++) {
		//get Answers
		dns_ans = get_resource(buf, &offset);
	}
	PPURPLE("Authoritative nameservers:");
	for(i =0; i < dns_head.arcount; i++) {
		//get Authoritative Nameservers
		dns_auth = get_resource(buf, &offset);
	}
	PPURPLE("Additional records:");
	for(i =0; i < dns_head.arcount; i++) {
		//get Additional Records
		dns_add = get_resource(buf, &offset);
	}
	return 0;
}


int main(int argc, char ** argv) {
int fd;
struct sockaddr_in dest_addr;
dns_question q;
	if (argc < 2) {
		printf("Not enough arguments. Please supply domain to lookup\n");
		exit(1);
	}
	PGREEN("My DNS Implementation");
	fd = socket( AF_INET, SOCK_DGRAM, 0 );
	if (fd <= 0) { perror("socket() error"); return 1; }

	dest_addr.sin_family = AF_INET;
	dest_addr.sin_port = htons(PORT);
	dest_addr.sin_addr.s_addr = inet_addr( NS );
	memset(&(dest_addr.sin_zero), '\0', 8);
	q.qtype = QTYPE_A;
	q.qclass = QCLASS_IN;
	q.qname = argv[1];
	make_query(fd, dest_addr, q);
	read_responce(fd, dest_addr);
	
	close(fd);
	return 0;
}