/*
 * Adapted from W.R. Stevens
 *   "Unix Network Programming", 1st ed. 
 */

#include  <stdio.h>
#include  <sys/types.h>
#include  <sys/socket.h>
#include  <netinet/in.h>
#include  <arpa/inet.h>

#include  <stdio.h>
#include  <strings.h> // bzero
#include  <string.h>  // strncpy
#include  <stdlib.h>  // exit
#include  <unistd.h>  // close

#include <stdint.h>   // uint8_t

int print_dns_name(char *, char *);

#define SERV_PORT   53
#define SERV_HOST_ADDR  "129.170.170.2"

// Pasted from Wireshark's Analyze > Follow > UDP stream > "C arrays" 
//  This represents the UDP DNS payload for an NS query for cs.dartmouth.edu
//  as generated by "dig ns cs.dartmouth.edu"
char req[] = {
     0xa8, 0x3c, 0x01, 0x00, 0x00, 0x01, 0x00, 0x00, 
     0x00, 0x00, 0x00, 0x00, 0x02, 0x63, 0x73, 0x09, 
     0x64, 0x61, 0x72, 0x74, 0x6d, 0x6f, 0x75, 0x74, 
     0x68, 0x03, 0x65, 0x64, 0x75, 0x00, 0x00, 0x02, 
     0x00, 0x01 };

#define MAX_MSG 1500

// DNS header with fixed-width fields
//   See http://unixwiz.net/techtips/iguide-kaminsky-dns-vuln.html for images.
typedef struct dns_resp_hdr {
     uint16_t tx_id;
     uint8_t  flags[2];
     uint16_t n_q;
     uint16_t n_a;
     uint16_t n_auth;
     uint16_t n_add;
} dns_resp_hdr_t;

int main(argc, argv)
int	argc;
char	*argv[];
{
     int sockfd;
     struct sockaddr_in	serv_addr;

     char buff[1500];
     int n, i; 

     uint16_t type, class, data_len;
     uint32_t ttl;

     dns_resp_hdr_t *r; 

     char *p;

     /*
      * Fill in the structure "serv_addr" with the address of the
      * server that we want to connect with.
      */
     
     bzero((char *) &serv_addr, sizeof(serv_addr));
     serv_addr.sin_family      = AF_INET;
     serv_addr.sin_addr.s_addr = inet_addr(SERV_HOST_ADDR);
     serv_addr.sin_port        = htons(SERV_PORT);
     
     /*
      * Open a UDP socket (a "datagram" socket, i.e., UDP; recall that TCP was SOCK_STREAM).
      */
     
     if ( (sockfd = socket(AF_INET, SOCK_DGRAM, 0)) < 0){
          perror("udp client: can't open DGRAM socket");
          exit(1);
     }

     /*
      * Connect to the server. Unlike TCP, no packets are sent; only kernel's
      *  data structures are filled in, so that next time you write/send to
      *  this socket, the IP/UDP packet knows where to go.
      */
     
     if (connect(sockfd, (struct sockaddr *) &serv_addr,
                 sizeof(serv_addr)) < 0){
          perror("udp client: can't connect to server");
          exit(1);
     }
     
     // Compare with sendto(), which does not need connect() above, because
     //   it specifies the sockaddr with every datagram.
     if( (n = write(sockfd, (void*) req, sizeof(req))) < 0 ){
          perror("udp client: write error");
          exit(1);
     }
     
     // Compare with recvfrom, which does not need a connect() either.
     n = recv(sockfd, (void*) buff, MAX_MSG, 0);
     if( n < 0 ){
          perror("Truncated response received.");
          goto out;
     }

     // This is the key to our parsing of the header: the struct gives us a "mask" for memory accesses.
     r = (dns_resp_hdr_t *) buff;
     printf( "Response length: %d\n", n );

     // Print packet in hex
     for (i = 0; i < n; i++) {
          printf("%02X%s", (uint8_t)buff[i], (i + 1)%16 ? " " : "\n");
     }
     printf("\n");

     // Printing out the fixed-length parts of the common DNS header:
     printf("tx_id 0x%x, flags %02X %02X, questions %d, answers %d, authorities %d, additional %d\n",
            ntohs(r->tx_id), r->flags[0], r->flags[1], ntohs(r->n_q), ntohs(r->n_a), ntohs(r->n_auth), ntohs(r->n_add) );
     
     p = buff + sizeof(dns_resp_hdr_t);

     // Print queries
     for( i = 0 ; i < ntohs(r->n_q); i++ ){
          // parse RR starting from name
          printf( "Query RR name: " );    
          p += print_dns_name( p, buff );
          
          // then print type and class
          type = ntohs( *(uint16_t *)p );
          p += 2;
          class = ntohs (*(uint16_t *)p );
          p += 2;
          printf( "  type %d, class %d\n", type, class );
     }  
     // Print answers 
     for( i = 0 ; i < ntohs(r->n_a); i++ ){
          // parse RR starting from name
          printf( "Answer RR name: " );    
          p += print_dns_name(p, buff);

          type = ntohs( *(uint16_t *)p );
          p += 2;
          class = ntohs (*(uint16_t *)p );
          p += 2;
          printf( "  type %d, class %d\n", type, class );

          ttl = ntohs( *(uint32_t *)p ); 
          p += 4;

          data_len = ntohs( *(uint16_t *)p ); 
          p += 2;
          printf( "  ttl %d, data_len %d ", ttl, data_len );

          // Data interpretation varies by type
          if( type == 1 && data_len == 4){
               printf( " %s\n", inet_ntoa( *(struct in_addr*) p ));
               p += 4;
          }
          else if( type == 2 ){
               print_dns_name( p, buff );
          } 
          else{
               printf( "WARN: don't know how to parse RR data for type %d length %d\n", type, data_len);
          }
          
          p += data_len;

     } // done printing answers
     for( i = 0 ; i < ntohs(r->n_auth); i++ ){
          printf( "Skipping authority RR: not implemented\n");
     }
     for( i = 0 ; i < ntohs(r->n_add); i++ ){
          printf( "Skipping additional RR: not implemented\n");
     }

out:
     close(sockfd);
     return 0;
}

/*
 * Args: pointer to a DNS-encoded name (or backreference), pointer to
 *   the start of the DNS payload (to resolve backreferences). 
 * Returns: the number of bytes occupied by the name, to skip over.
 */ 
int print_dns_name(char *p, char *buff)
{
     char *p0 = p;
     int8_t len  = *(uint8_t*)p;
     p++;

     // write out parts of the name from this record
     while( len > 0 ){
          fwrite( p, len, 1, stdout );
          printf(".");
          p += len; 
          len = *(uint8_t*)p;
          p++;
     }
     // at this point we either hit len==0 (the end of name)
     //  or a negative back-pointer

     if( len < 0 ){ // back-pointer, starts with 0xc0 
          uint16_t offset;
          offset = ntohs(*(uint16_t *)(p-1)) & ~0xc000;  // scan back one byte
          printf( "<backref offset %d>", offset );
          print_dns_name( buff+offset, buff ); // ignore return value
          printf("\n");
          p += 1; // skip offset
     }

     return p-p0;
}