#include <stdio.h>      /* standard C i/o facilities */
#include <stdlib.h>     /* needed for atoi() */
#include <unistd.h>  	/* defines STDIN_FILENO, system calls,etc */
#include <sys/types.h>  /* system data type definitions */
#include <sys/socket.h> /* socket specific definitions */
#include <netinet/in.h> /* INET constants and stuff */
#include <arpa/inet.h>  /* IP address conversion stuff */
#include <netdb.h>	/* gethostbyname */
#include <sys/time.h>


void tests(int,struct sockaddr_in *);
void filetransfer(int sk, struct sockaddr_in *s, char *filename);

/* client program:

   The following must passed in on the command line:
      hostname of the server (argv[1])
      port number of the server (argv[2])
*/



int main( int argc, char **argv ) {
  int sk;
  struct sockaddr_in server;
  struct hostent *hp;


  /* Make sure we have the right number of command line args */

  if (argc!=3) {
    printf("Usage: %s <server name> <port number>\n",argv[0]);
    exit(0);
  }

  /* create a socket 
     IP protocol family (PF_INET) 
     UDP (SOCK_DGRAM)
  */
  
  if ((sk = socket( PF_INET, SOCK_DGRAM, 0 )) < 0)
    {
      printf("Problem creating socket\n");
      exit(1);
    }

  server.sin_family = AF_INET;
  if ((hp = gethostbyname(argv[1]))==0) {
    printf("Invalid or unknown host\n");
    exit(1);
  }

  /* copy the IP address into the sockaddr 
     It is already in network byte order 
  */

  memcpy( &server.sin_addr.s_addr, hp->h_addr, hp->h_length);

  /* establish the server port number - we must use network byte order! */
  server.sin_port = htons(atoi(argv[2]));

  filetransfer(sk,&server,"medium");
  return(0);
}


/* code for building messages */

enum opcodes {RRQ=1,WRQ=2,DATA=3,ACK=4,ERR=5};

/* build rrq or wrq messages */
/* assumes that msg can hold at least 516 bytes! */
/* returns size of message or -1 */

int build_request(char *msg, unsigned short opcode, char *filename, char *mode) {

  if (strlen(filename)>100) {
    printf("File name too long\n");
    return(-1);
  }

  if (strlen(mode)>100) {
    printf("Transfer mode too long\n");
    return(-1);
  }

  if ((opcode != RRQ) && (opcode != WRQ)) {
    printf("build_request: invalid opcode (%d)\n",opcode);
    return(-1);
  }

  opcode=htons(opcode);
  memcpy(msg,&opcode,2);
  strcpy(msg+2,filename);
  strcpy(msg+2+strlen(filename)+1,mode);
  return(2+strlen(filename)+1+strlen(mode)+1);
}


/* build ack message */

int build_ack(char *msg, unsigned short block) {
  unsigned short opcode;

  opcode=htons(ACK);
  memcpy(msg,&opcode,2);
  block = htons(block);
  memcpy(msg+2,&block,2);
  return(4);
}

unsigned short get_shortval( char *p ) {
  unsigned short tmp;

  memcpy(&tmp,p,2);
  tmp = ntohs(tmp);
  return(tmp);
}

char *opcode_strings[] = {  "RRQ", "WRQ", "DATA", "ACK", "ERR" };

int get_opcode(char *msg) {
  unsigned short opcode;
  opcode = get_shortval(msg);
  return(opcode);
}


int get_blocknum(char *msg) {
  unsigned short blocknum;
  blocknum = get_shortval(msg+2);
  return(blocknum);
}

/* print out a TFTP message */
void print_msg( char *msg, int len ) {

  unsigned short opcode,blocknum,errcode;

  opcode = get_shortval(msg);

  if (opcode> ERR) {
    printf("Error - invalid opcode: %d\n",opcode);
    return;
  }

  printf("%s: ",opcode_strings[opcode-1]);

  switch(opcode) {
  case RRQ:
  case WRQ:
    printf("<%s> <%s>\n",msg+2,msg+2+strlen(msg+2)+1);
    break;
  case DATA:
    blocknum = get_shortval(msg+2);
    printf("BLOCK: %d (%d bytes)\n",blocknum,len-4); 
    break;
  case ACK:
    blocknum = get_shortval(msg+2);
    printf("BLOCK: %d\n",blocknum); 
    break;
  case ERR: 
    errcode = get_shortval(msg+2);
    printf("CODE: %d <%s>\n",errcode,msg+4);
    break;
  }

}


/*

void testit(int sk, struct sockaddr_in s) {
  char buff[1000];
  int len;
  int sent;
  int received;

  len = build_request(buff,RRQ,"small","octet");
  printf("SENDING ");
  print_msg(buff,len);

  if (len<0) {
    printf("Error, invalid request\n");
    return;
  }

  sent = sendto(sk,buff,len,NULL,&s,sizeof(s));
  if (sent!=len) {
    printf("Some kind of error sending!\n");
    return;
  }
  
  while (1) {
    len = sizeof(s);
    received = recvfrom(sk,buff,1000,NULL,&s,&len);
    if (received<0) {
      printf("Error calling recvfrom\n");
      return;
    }

    printf("RECEIVED: ");
    print_msg(buff,received);
  }
}
*/


int wait_for_response(int sk,struct sockaddr_in *s,char *buff, int max) {
  int len=max;
  fd_set rfds;
  struct timeval t = { 2,0 }; /* 2 seconds */ 
  int received;

  FD_ZERO(&rfds);
  FD_SET(sk,&rfds);
  if (select(sk+1,&rfds,NULL,NULL,&t)==0) {
    return(0);
  } else if (FD_ISSET(sk,&rfds)) {
    received = recvfrom(sk,buff,1000,0,(struct sockaddr*)s,&len);
    if (received<0) {
      printf("Error calling recvfrom\n");
      return(0);
    }
    printf("RECEIVED: ");
    print_msg(buff,received);
    return(received);
  }
  return(0);
}


void filetransfer(int sk, struct sockaddr_in *s, char *filename) {
  int sent;
  char buff[1000];
  int len;
  int opcode,block;
  int bnum=1;

  len = build_request(buff,RRQ,filename,"octet");
  printf("SENDING ");
  print_msg(buff,len);
  sent = sendto(sk,buff,len,NULL,(struct sockaddr*)s,sizeof(*s));

  while (1) {
    if ((len=wait_for_response(sk,s,buff,1000))==0) {
      printf("No Response\n");
      return;
    }
    opcode = get_opcode(buff);
    block = get_blocknum(buff);
    if ((opcode != DATA) || (block != bnum)) {
      printf("Invalid (not what I expected!)\n");
      return;
    }
    printf("Got block %d (%d)\n",block,len);

    build_ack(buff,block);
    sent = sendto(sk,buff,4,NULL,(struct sockaddr*)s,sizeof(*s));
    printf("SENDING ACK: ");
    print_msg(buff,4);
    if (len!=516) {
      break;
    }
    bnum++;
  }
}

void testmsg(int sk, struct sockaddr_in *s, char *msg, int len ) {
  int sent;
  char buff[1000];

  sent = sendto(sk,msg,len,0,(struct sockaddr*)s,sizeof(*s));
  if (sent!=len) {
    printf("Some kind of error sending (from sendto)!\n");
    return;
  }
  
  if (wait_for_response(sk,s,buff,1000)==0) {
    printf("No Response\n");
  }

}


char invalidopcode[]= "\x00\x08HiDave\x00";

char badrrq[] = "\x00\x01medium";

char rrq[] = "\x00\x01medium\x00octet\x00";

char rrq2[] = "\x00\x04\x00\x01";

char rrq3[] = "\x00\x02\x00\x01";

void tests(int sk, struct sockaddr_in *s) {

  struct sockaddr_in orig;

  memcpy(&orig,s,sizeof(struct sockaddr_in));


  /* send invalid opcode */
   printf("Sending invalid opcode\n");
  testmsg(sk,s,invalidopcode,sizeof(invalidopcode));


  /* send a RRQ with invalid filename (no null) */
  memcpy(s,&orig,sizeof(struct sockaddr_in));
  printf("Sending invalid RRQ\n");
  testmsg(sk,s,badrrq,sizeof(badrrq)-1); /* one byte too few - no null!*/

  /* send a valid RRQ */
  memcpy(s,&orig,sizeof(struct sockaddr_in));
  printf("Sending valid RRQ to get transfer started\n");
  testmsg(sk,s,rrq,sizeof(rrq)); 

  /* now send 1st ACK to new port */
  printf("Sending valid RRQ to get transfer started\n");
  testmsg(sk,s,rrq2,sizeof(rrq2)); 

  /* now send garbage to new port */
  printf("Sending Invalid stuff\n");
  testmsg(sk,s,rrq3,sizeof(rrq3)); 

}