#include "declarations.h"

/*
 * Global Variables
 * The Server socket descriptor : serverfd
 * Server address structure     : skaddr
 * Identification               : identification
 * Buffer to store the request  : buff
 * The size of the request      : msg_size
 */

int serverfd = -1;
struct sockaddr_in skaddr;
short identification;
unsigned char buff [MAX_MSG_SIZE]; 
static int msg_size = -1;

/*
 * Function Declarations
 */

/*
 * constructs the request from the hostname passed as a parameter. the
 * second parameter passed to this function is the buffer in which the
 * request should be stored.
 */
int construct_request (char *, char []);

/*
 * Skips the question field in the header of the response.
 * Arguments: 
 * Buff:	The buffer in which the response is stored
 * index: The pointer to the starting point in the buffer
 */
void skip_question( char *, int *);

/*
 * Prints out the addresses based on the last argument passed to this
 * function. The function iterates over resource records.
 * Arguments:
 * 1. char * buff : The buffer that contains the response
 * 2. int * index : The index at which reading should begin in the
 * buffer
 * 3. int numiter : The number of iterations/resource records
 * 4. int type    : The type of the query (0: Auth, 1: Ans, 2: Addrr)
 */
void output_addresses( char *, int *, int, int);

/*
 * Initializes a connection to the server. The code connects to a server
 * so that all error conditions will be communicated back to the client.
 * Argument: 1. The name of the server to connect to! This method is
 * called every time a new name server is returned so that the socket
 * can be connected to a new server!
 */
void connect_to_server(char* );

/*
 * High level Parser for the response from the name server. This method
 * calls skip_question, output_addresses and other response related
 * functions...
 * Arguments: 1. The response from the server in a buffer
 */
void parse_response( char [] );

/*
 * Handles the alarm conditions... Simply returns after printing a
 * message
 */
void signal_handler();

/*
 * The function that is the starting point for the recursion. If the
 * numanswers = 0 and additional rrs is > 0, this function is called.
 */
void callnextserver( char *, int *, int, int);


/*
 * Initializes the client by creating a socket to the server and
 * calling connect_to_server
 */
void init_client ( char * url) {
	//Create the socket
	
  /* create a socket 
     IP protocol family (PF_INET) 
     UDP protocol (SOCK_DGRAM)
  */
  
  //The value of protocol should be IPPROTO_UDP
	signal ( SIGALRM, signal_handler);
  if ( ( serverfd = socket( PF_INET, SOCK_DGRAM, 0 )) < 0) { 
    printf("Problem creating socket\n");
    exit(1);
  }

	connect_to_server( url);
}

void connect_to_server(char* url) {
  struct hostent *hostptr;

  /* 
	 * fill in an address structure that will be used to specify
     the address of the server we want to connect to

     address family is IP  (AF_INET)

     server IP address is found by calling gethostbyname with the
     name of the server (entered on the command line)
	 *
   */

  //You need this because the address needs to be put into the socket
	bzero( &skaddr, sizeof( skaddr));
  skaddr.sin_family = AF_INET;

	alarm( WAIT_TIME);
	if ( (hostptr = gethostbyname( url)) == NULL) { 
		//somehow, this doesn't give an error. The error occurs only in a
		//read!
		fprintf(stderr, "Error getting the host name... Using the ip address!\n");
		if( ( skaddr.sin_addr.s_addr = inet_addr(url)) == -1) {
			fprintf(stderr, "Invalid IP address: %s\n",url);
			alarm(0);
		  exit ( 1);
		} 
	} 
	else {
		//The gethostbyname function worked!
		memcpy( &skaddr.sin_addr.s_addr, hostptr->h_addr, hostptr->h_length);
	}
	alarm (0);

	//memcpy( &skaddr.sin_addr.s_addr, hostptr->h_addr, hostptr->h_length); 
	//memcpy avoids the memory leads that can otherwise occur
	free( hostptr);
	
  skaddr.sin_port = htons( DNS_PORT);

	/*
	 * We have now created the socket. All we need to now do is write to
	 * the server and see what happens!
	 */
	//Addendum: One should connect the UDP socket for two reasons:
	//1. It gives better performace (partridge and park)
	//2. ICMP port unreachable errors are returned to the socket if it is
	//connected. This gives a better picture of what exactly occured when
	//the serverdid not return a message... 
	//3. Even though we are connecting the socket, an alarm should be
	//implemented so that the recvfrom does not halt because of a lost
	//reply from the server.
	
	if( connect( serverfd, (struct sockaddr *) &skaddr, sizeof( skaddr)) == -1) {
			printf("An error occured ... \n");
			printf ( "The error was : \n%s\n", strerror( errno));
	}

	return;
}

void retrieve_address( char * hostname) {
	unsigned char read_buff [MAX_MSG_SIZE];
  int num_read;
  int i;

	bzero( read_buff, sizeof( read_buff));
	//The socket has been created...
	//Now just check if the socket is not equal to -1 and if it is
	//connected to a valid server, in short, if init_client has been
	//called
	
	if( serverfd == -1 || skaddr.sin_port != htons( DNS_PORT)) {
		printf( "You need to call init_client before calling this function... Bummer!\n");
		return;
	}

	//Talk to the server now... Send the request!
 	if( (hostname != NULL) && (msg_size == -1)) 
		msg_size = construct_request( hostname, buff);

	//send request
	//if( (i = sendto( serverfd, buff, msg_size, 0, (struct sockaddr *) &skaddr, sizeof( skaddr))) < 0) {
	if( (i = write( serverfd, buff, msg_size)) < 0) {
		printf( "Error occured!\n"); 
		printf( "%s\n", strerror( errno));
		exit( 1);
	}
	
	printf( "%d = num written\n", i);
	//read response
	for( i = 0; i < NUM_RETRIES; i++) {
		//get the buff now!
		
	  signal ( SIGALRM, signal_handler);
		alarm(WAIT_TIME);
		num_read = read ( serverfd, read_buff, MAX_MSG_SIZE);
		//num_read = recvfrom ( serverfd, read_buff, MAX_MSG_SIZE, 0, NULL, NULL);
		if(num_read < 0) {
			if( errno == EINTR) {
				perror( "Error in read... I think server timed out!\n");
				printf ( "%s\n", strerror ( errno));
				printf( " Retrying...%d \n", i);
				alarm(0);
				continue;
			} else {
				printf( "Error in read... Exiting\n");
				exit( 1);
			}
		} else {
			printf( "%d = num bytes read\n", num_read);
			alarm(0);
			break;
		}
		
		if(num_read > MAX_MSG_SIZE) {
			perror( "The size of the response is greater than 512.\nI am not going to use TCP to get a response... \n");
			exit( 1);
		}
	}

	if ( i == NUM_RETRIES) {
		printf( "Retried too many times. Server not responding...\n");
		exit( 1);
	}

	/*
	for( i = 0; i < num_read; i++) {
		putchar( read_buff[i]);
	}
	putchar( '\n');
	*/

	parse_response( read_buff);
	//print out response
}

//char *  construct_request (char * hostname) {
//This will not work... I will return the value of a local variable and
//that will not reach the calling function (the stack is killed!!)

int  construct_request (char * hostname, char buff[]) {
  char * tmp;
	int length;
	int currentIndex = 0;
	unsigned char c = 0x0;
	int i = 0;
	unsigned short flag = 1;
	short htons1 = htons(1); //, htons = htons(0);
	short temp = 0;
	
	//Initialize the message to be all zeros
	bzero( buff, sizeof( buff));
	
	//Create the id
	identification = strlen( hostname);

	temp = htons(identification);
	memcpy( buff + currentIndex, & temp, 2);
	
	currentIndex += 2;	
	
	//Create the flag!
	//If I set this, the server does the recursion for me. I
	//am going to send you two versions...
	flag <<= 8; 
	//printf("%d = flag\n", flag);

	temp = htons(flag);
	memcpy( buff + currentIndex, &temp, 2);
	currentIndex += 2;
	
	//Question Count
	memcpy ( buff + currentIndex,  &htons1, 2);
	currentIndex += 2;
	
	//Answer Count
	buff[currentIndex ++] = 0;
	buff[currentIndex ++] = 0;
	
	//Authority RECS
	buff[currentIndex ++] = 0;
	buff[currentIndex ++] = 0;
	
	//Additional RRs
	buff[currentIndex ++] = 0;
	buff[currentIndex ++] = 0;
	
	/*
	 * Done with the HEADER. Now to create the QUESTION!
	 */

	//Encode the address:
	//The address is of the form www.cs.rpi.edu
	//Assign currentIndex to 2
	tmp = strtok( hostname, ".");
	
	do {
		length = strlen( tmp);
		
		if (length > 63)  {
			printf (" you have requested the url of an invalid hostname \n<label length > 63> Exiting.... error code 20001\n");
			exit( 1);
		}
		
		//Bit 7 and 6 must always be 0
		c = (unsigned char) length;
		
		buff[currentIndex++] = c;
		
		for (i = 0; i < length && currentIndex < MAX_MSG_SIZE; i++, currentIndex++) 
		{
			buff[currentIndex] = tmp[i];
		}
		
		if( currentIndex > MAX_MSG_SIZE) {
			fprintf( stderr, 
					" The address you have provided is way too large for a dns message buffer. Exiting!\n");
			exit( 1 );
		}
	} while ( (tmp = strtok( NULL, ".")) && (currentIndex < MAX_MSG_SIZE) );
	
	buff[currentIndex ++] = (unsigned char) 0;

	//The address has been encoded
	
	//Query Type
	memcpy ( buff + currentIndex, &htons1, 2);
	currentIndex += 2;
	
	//Query Class
	memcpy ( buff + currentIndex, &htons1, 2);
	currentIndex += 2;

	/*for( i = 0; i < currentIndex; i++) {
		putchar( buff[i]);
	}*/
	printf( "\n%d = num bytes in the request!\n", currentIndex);
	return currentIndex;
}	

void parse_response( char response[]) {
	//First get the header: Size 12 bytes
	char header [12];
  int currentIndex = 0;
	int insidehdr = 0;  
	unsigned short flag = 0;
	short numanswers = 0;
	short authrecs = 0, additionalrr = 0;
	unsigned short temp = 0;
	unsigned short rcodemask = 0xf;
	
	bzero( header, sizeof( header));
	memcpy( header, response, 12);
	
	//parse the response header but first, update the currentIndex because
	//it will not be used till the header has been parsed!
	currentIndex += 12;
	
	memcpy( &temp, header + insidehdr, 2);
	temp = ntohs(temp);

	if (temp != identification) {
		printf("A stray packet sent. \nIdentification of the response does not match that of the request.\nYou guys cannot fool me!\n");
		exit( 1);
	}

	insidehdr += 2;
	//copy the flag from the header into flag
	memcpy( &flag, header + insidehdr, 2);
	flag = ntohs( flag);
	printf( "%x = flag\n", flag);

  if( flag & rcodemask) {
		printf( "RCODE error... Exiting!\n");
		exit( 1);
	}
	
	insidehdr +=4; //skip over the num questions

	//copy the number of answers from the header to numanswers
	memcpy( &numanswers, header + insidehdr, 2);
	numanswers = ntohs( numanswers);
	printf( "%d = numanswers\n", numanswers);
	insidehdr += 2;
	
	printf( "Num answers = %d\n", numanswers);

	//num authority records
	memcpy( &authrecs, header + insidehdr, 2);
	insidehdr += 2;
	authrecs = ntohs( authrecs);
	printf( "Num auth = %d\n", authrecs);

	//num additional resource records
	memcpy( &additionalrr, header + insidehdr, 2);
	insidehdr += 2;
	additionalrr = ntohs( additionalrr);
	printf( "Num additional = %d\n", additionalrr);

	//the number of answers is not 0, so lets print out the answers got!
	skip_question( response, &currentIndex);
	
	//Get the answer records
	if( !numanswers ) {
		
		//There can be some additional rrs giving me servers to contact
		if( additionalrr)
			//This function will make the recursive call!
			callnextserver( response, &currentIndex, authrecs, additionalrr); 

		//if there are no additional rrs and no answers sent in the
		//response, I should return to the point where this function was
		//called
		return;
	} 
	
	//this is called only if the number of answers is not equal to 0. If
	//this is the case, the answer must be printed out. Call type = 1
	output_addresses( response, &currentIndex, numanswers, 1);
}

void skip_question( char * response, int * index) {
	//Skips a question....
	unsigned char value = 0;
	int i = 0;
	
	//The index points to the first element of the question
	do {
		value = response[*index];
		*index = *index + 1;		//next byte after value...
		for( i = 0; i < value; i++) {
			putchar( response[ *index]);
			*index = *index + 1; //next byte... 
		}
		putchar('.');
	} while ( value); //will continue till value != 0
	putchar('\n');
	*index += 4; //skips over the query type and class
}

//There can three types of callTypes
//0. authority rec, ignore
//1. address, print it
//2. additional rr, recurse

void output_addresses( char * response, int * index, int iterations, int callType) {
	int i = 0;
	short dlen = 0;
	short atype, aclass;
	unsigned char temp = 0;
	struct 	in_addr addr;
	char * address = NULL;
	
	printf("%d = iterations, %d = callType\n", iterations, callType);
	for( i = 0; i < iterations; i++) {
		if (callType) 
			printf("----------answer or additional record... printing-------\n");
		else 
			printf("----------authoritative record... skipping-----------\n");
	
		//determine the nature of the first byte of the rr
		temp = response[*index];

		//if the first two bits of temp == 11, this is a pointer field,
		//increment index by 2 else, this is a variable length domain name
		//with the first two bits as 00 (so that the max length of a label
		//is 63)
		temp >>= 6; // 8 - 6 = 2 if temp is not equal to 0, this is a pointer
		
		if( temp == 3) {
			//This is a pointer, increment index by 2
			*index += 2;
		} else 
			if ( temp == 0) {
				skip_question( response, index);
				*index -= 4; //need to know the type and class
			} else {
				//Anything else... this implies an unused type... Exit
				printf(" The response contains a currently unassigned value... exiting!");
				exit( 1);
			}

		//type
		memcpy ( &atype, response + *index, 2);
		*index += 2;
		if (callType == 2) 
			printf(" Type : %d\n", atype);

		//class
		memcpy ( &aclass, response + *index, 2);
		*index += 2;
		if (callType == 2) 
			printf(" Class : %d\n", aclass);
	
		*index += 4; //time to live skipped!

		//length of data
		memcpy( &dlen, (response + *index), 2);
		*index += 2;
		dlen = ntohs( dlen);
		
		if(atype != 1) {
			//not an answer! continue
			*index += dlen;
			printf(" ---------------------------------------------------\n");
			continue;
		}
		
		//An answer. Check if this was called with type 1 or 2
		if(atype == 1 && (callType > 0)) {
			memcpy( &addr.s_addr, response + *index, dlen);
			address = inet_ntoa(addr);
			
			*index += dlen;
			//switch here...
			if( callType == 1)
				printf( "IP Address is: %s\n", address); 
			else if( callType == 2) {
				
				/* ********************************************************
				 *                    THE RECURSION                       *
				 * ********************************************************/
				//1. connecttoserver(with the server name)
				//2. retrieve_address( with some junk value!)
				
				printf( "Next DNS Server IP Address is: %s\n", address); 
				printf(" ---------------------------------------------------\n");
				
				connect_to_server( address);
				
				//NULL because you can pass any junk value, the request is not going 
				//to be recreated
				retrieve_address( NULL); 
			}
			//We have ignored the case when the value is 0
			printf(" ---------------------------------------------------\n");
			//break;
		}
	}
}

void callnextserver(char response[], int *index, int authrecs, int addrr) {
	printf( "---------------------Call next server ---------------------\n");
	
	//Check if authrecs are greater than 0!
	if( authrecs) 
		output_addresses( response, index, authrecs, 0);

	//additionalrr is always greater than 0 if this call is made!
	output_addresses( response, index, addrr, 2);
}
	
void signal_handler() {
	printf ( "A system call timed out... \n");
}

int main(int argc, char ** argv) {
	
  if ( argc != 3) {
    printf( "Usage: dnsclient <server> <hostname>\n");
    exit(1);
  }
	
	init_client(argv[1]);
	if( strlen( argv[2]) > MAX_MSG_SIZE) {
			printf( "Trying to overflow my buffers, eh?! Not as dumb as you make me out to be!\n"); 
			exit( 1);
	}
	retrieve_address( argv[2]);
	return 0;
}