/**************************************************
* TFTP client compatible with RFC-1350 
* compile under visiual c++ or borland c++
***************************************************/


#define _VC /* if compile under visiual c++ else undefine this*/
#include <stdio.h>
#include <winsock.h>
#include <conio.h>

#ifndef MAKEWORD
#define MAKEWORD(l,h) ((WORD)(((BYTE)(l))|(((WORD)(BYTE)(h))<<8)))
#endif
#define WSA_MAJOR_VERSION 1
#define WSA_MINOR_VERSION 1
#define WSA_VERSION MAKEWORD(WSA_MAJOR_VERSION, WSA_MINOR_VERSION)

/* read/write request packet format
     2 bytes     string    1 byte     string   1 byte
     ------------------------------------------------
    | Opcode |  Filename  |   0  |    Mode    |   0  |
     ------------------------------------------------

*/
#define TFTP_RRQ 1   /*Read request (RRQ)*/
#define TFTP_WRQ 2   /*Write request (WRQ) */

/* DATA packet format
       2 bytes     2 bytes      n bytes
      ----------------------------------
     | Opcode |   Block #  |   Data     |
      ----------------------------------
*/
#define TFTP_DATA 3  /*Data (DATA)*/

/* ACK packet format
      2 bytes     2 bytes
      ---------------------
     | Opcode |   Block #  |
      ---------------------
*/
#define TFTP_ACK 4   /*Acknowledgment (ACK)*/

/*ERROR packet format
    2 bytes     2 bytes      string    1 byte
    -----------------------------------------
   | Opcode |  ErrorCode |   ErrMsg   |   0  |
    -----------------------------------------
*/
#define TFTP_ERROR 5 /*Error (ERROR)*/

#define TFTP_NETASCII 0
#define TFTP_OCTET 1
#define TFTP_WSTAT_FIRSTACK 0
#define TFTP_WSTAT_NEXTACK 1
#define TFTP_WSTAT_LASTACK 2
#define MAX_RETRY 3
#define TFTP_NOTEND_DATALEN 512+2+2

#ifdef _VC
#pragma comment( lib, "Wsock32.lib" )
#endif

typedef void (* CMDFUNC)(char [][256],int pcount);
typedef struct _cmdnum{
 char *cmd;
 int num;
 int paramcount;
 CMDFUNC callback;
}CMDNUM,*PCMDNUM;

void connectto(char cmd[][256],int pcount);
void setoctet(char cmd[][256],int pcount);
void setascii(char cmd[][256],int pcount);
void quit(char cmd[][256],int pcount);
void showhelp(char cmd[][256],int pcount);
void test(char cmd[][256],int pcount);
void getfile(char cmd[][256],int pcount);
void putfile(char cmd[][256],int pcount);
int stripcmd(char *s,char cmd[][256]);
void parsecmd(char *s);
int  getcmdnum(char *s);
int makereq(char type,int mode,char *filename,char *buffer,int size);
int makeack(unsigned short num,char *buffer,int size );
void showsysinfo();

CMDNUM cmdlist[] = {
 {"help",1,0,showhelp},
 {"exit",2,0,quit},
 {"test",3,0,test},
 {"get",4,1,getfile},
 {"put",5,2,putfile},
 {"octet",6,0,setoctet},
 {"ascii",7,0,setascii},
 {"connect",8,1,connectto}};

char *helptext = "help: show this text\n\
exit: exit pragram\n\
ctet: set file mode to octet\n\
ascii: set file mode to netascii\n\
connect remoteip: connect to server\n\
get filename: get file from server\n\
put localname remotefilename: upload file to server\n";

SOCKET sock = INVALID_SOCKET;
char desthost[256] = "127.0.0.1"; 
int filemode = TFTP_NETASCII;
int main(int argc, char* argv[])
{

 char cmd[256];
 WSADATA stWSAData;
    int ret = 0; 
 sockaddr_in addr;
 showsysinfo();
 if(WSAStartup(WSA_VERSION, &stWSAData)!=0)
 {
  printf("Can't start Socket \n");
  exit(0);
 }
    
    sock = socket(PF_INET,SOCK_DGRAM,0);
 if(sock==INVALID_SOCKET)
 {
  printf("Can't create socket \n");
  exit(0);
 }
 addr.sin_family = PF_INET;
 addr.sin_port = INADDR_ANY;
 addr.sin_addr.s_addr  = INADDR_ANY;
 if(bind(sock,(struct sockaddr *)&addr,sizeof(addr))!=0)
 {
      printf("Can't bind socket \n");
   exit(0);
 }
 while(true)
 {
   fflush( stdin );
   printf("#");
   gets(cmd);
   parsecmd(cmd);
 }

 return 0;
}

void showsysinfo()
{
  printf("TFTP client version 1.0\n");

}
int stripcmd(char *s,char cmd[][256])
{
  int i=0;
  char *token=NULL;
  char seps[] = " ,\t\n";
  token = strtok( s, seps );
  while(token!=NULL)
  {
     if (i>2) break; 
     strcpy(cmd[i],token);
  token = strtok( NULL, seps );
  i++;
  }
  return i; 
}
int  getcmdnum(char *s)
{
  int i = 0;
  for(i=0;i<sizeof(cmdlist)/sizeof(CMDNUM);i++)
  {
    if(stricmp(s,cmdlist[i].cmd)==0)
 {
   return i;
 }
  }
  return -1;  
}

void parsecmd(char *s)
{
   char cmd[3][256];
   int pcount = 0;
   int num = -1;
   pcount = stripcmd(s,cmd);
   num=getcmdnum(cmd[0]);
   if(num==-1)
   {
     printf("No such commond \n");
  return;
   }
   else
   {
     cmdlist[num].callback(cmd,pcount-1);
   }
}

void quit(char cmd[][256],int pcount)
{
 printf("exit to system \n");
 closesocket(sock);
 exit(0);
}

void showhelp(char cmd[][256],int pcount)
{
   printf(helptext);
}

void test(char cmd[][256],int pcount)
{
}

void setoctet(char cmd[][256],int pcount)
{
  filemode = TFTP_OCTET;
  printf("Set file mode to octet\n");
}

void setascii(char cmd[][256],int pcount)
{
  filemode = TFTP_NETASCII;
  printf("Set file mode to netascii\n");
}

void connectto(char cmd[][256],int pcount)
{
  if(pcount<1)
  {
 printf("usage: connect remoteip \n");
 return;
  }
  strcpy(desthost,cmd[1]);
}

int makeack(unsigned short num,char *buffer,int size )
{
  int pos = 0;
  buffer[pos] = 0;
  pos++;
  buffer[pos] = TFTP_ACK;
  pos++;
  buffer[pos] = (char)(num>>8);
  pos++;
  buffer[pos] = (char)num;
  pos++;
  return pos;
}

int makereq(char type,int mode,char *filename,char *buffer,int size)
{
  int pos = 0;
  unsigned int i = 0;
  char s[32] = "";
  if(mode==TFTP_NETASCII)
    strcpy(s,"netascii");
  else
 strcpy(s,"octet");
  buffer[pos] = 0;
  pos++;
  buffer[pos] = type;
  pos++;
  for(i=0;i<strlen(filename);i++)
  {
    buffer[pos] = filename[i];
    pos++;
  }
  buffer[pos] = 0;
  pos++;
  for(i=0;i<strlen(s);i++)
  {
    buffer[pos] = s[i];
    pos++;
  }
  buffer[pos] = 0;
  pos++;
  return pos;
}

int makedata(int num,char *data,int datasize,char *buffer,int bufsize)
{
  int pos = 0;
  buffer[pos] = 0;
  pos++;
  buffer[pos] = TFTP_DATA;
  pos++;
  buffer[pos] = (char)(num>>8);
  pos++;
  buffer[pos] = (char)num;
  pos++;
  memcpy(&buffer[pos],data,datasize);
  pos = pos + datasize;
  return pos;

}

void getfile(char cmd[][256],int pcount)
{
  char sendbuf[1024] = {0};
  char recvbuf[1024] = {0};
  sockaddr_in addr;
  sockaddr_in from;
  int fromlen = 0;
  int ret = 0;
  int len = 0 ;
  fd_set  fdr;
  int retry = 0;
  struct  timeval timeout = {5,0};
  int stat = 0;
  int lastdata = 0;
  FILE *file;
  int flen = 0;
  int c;
  if(pcount!=1)
  {
   printf("usage: get filename\n");
   return;
  }
  if((file=fopen(cmd[1],"r"))!=NULL)
  {
    printf("File %s already exits,overwrite? y/n ",cmd[1]);
 while(true)
 {
   c = getch();
   if('Y'==toupper(c))
   {
        printf("\n");
  fclose(file);
  break;
   }
   else if('N'==toupper(c))
   {
        printf("\n");
  fclose(file);
  return;
   }
 }
  }
  if((file=fopen(cmd[1],"w+b"))==NULL)
  {
    printf("Can't create file\n");
 return;
  }
  len = makereq(TFTP_RRQ,filemode,cmd[1],sendbuf,sizeof(sendbuf));
  addr.sin_family =PF_INET;
  from.sin_family =PF_INET;
  addr.sin_port = htons(69);
  addr.sin_addr.s_addr   = inet_addr(desthost);
  ret = sendto(sock,sendbuf,len,0,(sockaddr *)&addr,sizeof(addr));
  while(true)
  {
    FD_ZERO(&fdr);
    FD_SET(sock, &fdr);
 ret = select(sock, &fdr, NULL,NULL, &timeout);
 if(SOCKET_ERROR==ret)
 {
      printf("Socket error \n");
   fclose(file);
      return;
 }
 else if(0==ret)
 {
      if(MAX_RETRY==retry)
   {
     printf("Time Out \n");
  fclose(file);
  return;
   }
      sendto(sock,sendbuf,len,0,(sockaddr *)&addr,sizeof(addr));
      retry++;
 }
 else
 {
   if (FD_ISSET(sock,&fdr))
   {
        retry = 0;
     fromlen = sizeof(sockaddr);
        ret = recvfrom(sock,recvbuf,sizeof(recvbuf),0,(sockaddr *)&from,&fromlen);
     if(TFTP_ERROR==recvbuf[1])
  {
          fclose(file);
          printf("Error %d: %s \n",recvbuf[3],&recvbuf[4]);
    return;
  }
     if(0==stat)
  {
          addr.sin_port = from.sin_port ;
    stat = 1;
  }
        if(TFTP_DATA==recvbuf[1])
  {
          lastdata = recvbuf[2]*256 + recvbuf[3];
    len = makeack(lastdata,sendbuf,sizeof(sendbuf));
          sendto(sock,sendbuf,len,0,(sockaddr *)&addr,sizeof(addr));
    if(ret<TFTP_NOTEND_DATALEN)
    {
   fwrite(&recvbuf[4],1,ret-4,file);
   flen = flen + ret -4;
            fclose(file);
   printf("total %d byte received\n",flen);
   return;
    }
          else
    {
            fwrite(&recvbuf[4],1,512,file);
   flen = flen + 512;
   printf("%d byte received\r",flen);
    }
  }
   }
 }

  }
}

void putfile(char cmd[][256],int pcount)
{
  char sendbuf[1024] = {0};
  char recvbuf[1024] = {0};
  char databuf[1024] = {0};
  sockaddr_in addr;
  sockaddr_in from;
  int fromlen = 0;
  int ret = 0;
  int len = 0 ;
  fd_set  fdr;
  int retry = 0;
  struct  timeval timeout = {5,0};
  int stat = TFTP_WSTAT_FIRSTACK;
  int lastack= 0;
  FILE *file;
  int flen = 0;
  int blocknum = 0;
  size_t rlen = 0;
  if(pcount!=2)
  {
   printf("usage: put localfilename remotefilename \n");
   return;
  }
  if((file=fopen(cmd[1],"r"))==NULL)
  {
    printf("File %s not found \n",cmd[1]);
 return;
  }
  len = makereq(TFTP_WRQ,filemode,cmd[2],sendbuf,sizeof(sendbuf));
  addr.sin_family =PF_INET;
  addr.sin_port = htons(69);
  addr.sin_addr.s_addr   = inet_addr(desthost);
  ret = sendto(sock,sendbuf,len,0,(sockaddr *)&addr,sizeof(addr));
  if((file=fopen(cmd[1],"r"))==NULL)
  {
    printf("Can't Open file %s\n",cmd[1]);
 return;
  }
  while(true)
  {
    FD_ZERO(&fdr);
    FD_SET(sock, &fdr);
 ret = select(sock, &fdr, NULL,NULL, &timeout);
 if(SOCKET_ERROR==ret)
 {
      printf("Socket error \n");
   fclose(file);
      return;
 }
 else if(0==ret)
 {
      if(MAX_RETRY==retry)
   {
     printf("Time Out \n");
  fclose(file);
  return;
   }
      sendto(sock,sendbuf,len,0,(sockaddr *)&addr,sizeof(addr));
      retry++;
 }
 else
 {
      retry = 0;
   fromlen = sizeof(sockaddr);
      ret = recvfrom(sock,recvbuf,sizeof(recvbuf),0,(sockaddr *)&from,&fromlen);
      if(TFTP_ERROR==recvbuf[1])
   {
        fclose(file);
        printf("Error %d: %s \n",recvbuf[3],&recvbuf[4]);
     return;
   }
      if(TFTP_ACK==recvbuf[1])
   {
        lastack = recvbuf[2]*256 + recvbuf[3];
     switch(stat)
  {
       case TFTP_WSTAT_FIRSTACK:
            if(0==lastack)
   {
              stat = TFTP_WSTAT_NEXTACK;
              addr.sin_port = from.sin_port ;
     rlen = fread(databuf,1,512,file);
     flen = flen + rlen;
           if(rlen<512 && feof(file))
     {
                stat = TFTP_WSTAT_LASTACK;
     }
           else if(ferror(file))
     {
                printf("Error: read file\n");
          fclose(file);
          return;
     }
              blocknum++;
     len = makedata(blocknum,databuf,rlen,sendbuf,sizeof(sendbuf));
        sendto(sock,sendbuf,len,0,(sockaddr *)&addr,sizeof(addr));
     printf("%d byte send\r",flen);
   }
      else
   {
     fclose(file);
     printf("Error Ack Number");
     return;
   }
      break;
    case TFTP_WSTAT_NEXTACK:
            if(lastack==blocknum)
   {
              rlen = fread(databuf,1,512,file);
     flen = flen + rlen;
           if(rlen<512 && feof(file))
     {
                stat = TFTP_WSTAT_LASTACK;
     }
           else if(ferror(file))
     {
                printf("Error: read file\n");
          fclose(file);
          return;
     }
              blocknum++;
     len = makedata(blocknum,databuf,rlen,sendbuf,sizeof(sendbuf));
        sendto(sock,sendbuf,len,0,(sockaddr *)&addr,sizeof(addr));
     printf("%d byte send\r",flen);
   }
      else
   {
     fclose(file);
     printf("Error Ack Number");
     return;
   }
      break;
    case TFTP_WSTAT_LASTACK:
      if(lastack==blocknum)
   {
     printf("%d byte send\n",flen);
     return;
   }
      else
   {
     fclose(file);
     printf("Error Ack Number");
       return;
   }
      break;
  }
   }
 }
  }
 }