/*
 * Example of client using TCP protocol, adapted from
 *   "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_TCP_PORT   53
#define SERV_HOST_ADDR  "129.170.170.2"

// Pasted from Wireshark's "Analyze > Follow > TCP stream",
//  generated by "dig +tcp a cs60.cs.dartmouth.edu"
// Note: unlike a UDP payload, this one has a two-byte
//  length field up front, 0x00 0x27 = 39 bytes after
//  the length field, 41 bytes all told.
char req[] = { 
     0x00, 0x27, 0x1d, 0x62, 0x01, 0x00, 0x00, 0x01, 
     0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04, 0x63, 
     0x73, 0x36, 0x30, 0x02, 0x63, 0x73, 0x09, 0x64, 
     0x61, 0x72, 0x74, 0x6d, 0x6f, 0x75, 0x74, 0x68, 
     0x03, 0x65, 0x64, 0x75, 0x00, 0x00, 0x01, 0x00, 
     0x01 };

#define MAX_MSG 1500

typedef struct tcp_dns_resp_hdr {
     uint16_t len;    // not there in UDP payloads! 
     uint16_t tx_id;
     uint8_t  flags[2];
     uint16_t n_q;
     uint16_t n_a;
     uint16_t n_auth;
     uint16_t n_add;
} tcp_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;

     tcp_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_TCP_PORT);
     
     /*
      * Open a TCP socket (an Internet stream socket).
      */
     
     if ( (sockfd = socket(AF_INET, SOCK_STREAM, 0)) < 0){
          perror("tcp client: can't open stream socket");
          exit(1);
     }

     /*
      * Connect to the server.
      */
     
     if (connect(sockfd, (struct sockaddr *) &serv_addr,
                 sizeof(serv_addr)) < 0){
          perror("tcp client: can't connect to server");
          exit(1);
     }
     
     //send(sockfd, (void*) msg, MSG_LEN, 0);
     write(sockfd, (void*) req, sizeof(req));
     
     n = recv(sockfd, (void*) buff, MAX_MSG, 0);
     if( n < 0 ){
          perror("Truncated response received.");
          goto out;
     }

     r = (tcp_dns_resp_hdr_t *) buff;
     printf( "Response length: %d\n", ntohs(r->len) );

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

     // Print 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(tcp_dns_resp_hdr_t);

     for( i = 0 ; i < ntohs(r->n_q); i++ ){
          // parse RR starting from name
          printf( "Query RR name: " );    
          p += print_dns_name( p, buff+2 );
          
          // 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+2);

          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;
}