/*
 * bsk_tftpd.cpp
 *
 *  Created on: 25/ago/2016
 *      Author: chessaa
 */

#include "bsk_tftpd.h"

#include <string.h>
#include <stdlib.h>

#include <stdio.h>

/*
 * Trivial File Transfer Protocol (IEN-133)
 */
#define TFTP_SEGSIZE          512     /* data segment size */
//#define TFTP_MAXSEGSIZE     16384       /* data segment size */
#define TFTP_MAXSEGSIZE     1428  /*16384 */      /* data segment size */
#define TFTP_MINSEGSIZE         8       /* data segment size */
/*
 * Packet types.
 */
#define TFTP_RRQ    01          /* read request */
#define TFTP_WRQ    02          /* write request */
#define TFTP_DATA   03          /* data packet */
#define TFTP_ACK    04          /* acknowledgement */
#define TFTP_ERROR  05          /* error code */
#define TFTP_OACK   06          /* option acknowledgement */


typedef struct {
  unsigned char  th_opcode[2];      /* packet type */
  union {
    unsigned char tu_block[2];   /* block # */
    unsigned char tu_code[2];    /* error code */
    unsigned char tu_stuff[1];    /* request packet stuff */
  } th_u;
} tftphdr_t;


typedef struct
{
  tftphdr_t header;
  char    th_data[2];     /* data or error string */
} tftppkt_t;


static void tftp_htons(unsigned char* const p, unsigned short v)
{
  p[0]=(unsigned char)(v>>8);
  p[1]=(unsigned char)(v);
}

static unsigned short tftp_ntos(const unsigned char* n)
{
  return n[1]|(n[0]<<8);
}

#define  th_block    th_u.tu_block
#define  th_code     th_u.tu_code
#define  th_stuff    th_u.tu_stuff
#define  th_msg      th_data

#define TFTP_DATA_HEADERSIZE ( offsetof (struct tftphdr, th_data ) )
#define TFTP_ACK_HEADERSIZE  (  offsetof (struct tftphdr, th_block )  \
                            + sizeof ( ((struct tftphdr *) (0))->th_block) )

/*
 * Error codes.
 */
#define  EUNDEF      0       /* not defined */
#define  ENOTFOUND   1       /* file not found */
#define  EACCESS     2       /* access violation */
#define  ENOSPACE    3       /* disk full or allocation exceeded */
#define  EBADOP      4       /* illegal TFTP operation */
#define  EBADID      5       /* unknown transfer ID */
#define  EEXISTS     6       /* file already exists */
#define  ENOUSER     7       /* no such user */
#define  EBADOPTION  8        /* bad option */
#define  ECANCELLED 99      /* cancelled by administrator */


/*
 * options
 */
#define TFTP_OPT_TSIZE     "tsize"
#define TFTP_OPT_TIMEOUT   "timeout"
#define TFTP_OPT_BLKSIZE   "blksize"
#define TFTP_OPT_MCAST     "multicast"
#define TFTP_OPT_PORT      "udpport"

#define IS_OPT(s,opt)   (lstrcmpi (s, opt)==0)

#define PKTSIZE             TFTP_SEGSIZE+4
#define MAXPKTSIZE          TFTP_MAXSEGSIZE+4

static int fake_log_delegate(bsk_tftpd_context_id_t, const char* msg, ...)
{
  (void)msg;
  return 0;
}

static int fake_send_delegate(bsk_tftpd_context_id_t, unsigned int, unsigned int, const void* const data, unsigned int size)
{
  (void)data;
  (void)size;
  return 0;
}

static bsk_tftpd_receive_status_t fake_receive_delegate(bsk_tftpd_context_id_t, bsk_tftpd_receive_delegate_operation_t op, const void* const data, unsigned int block, unsigned int offset, unsigned int size_bytes)
{
  static bsk_tftpd_receive_status_t reply={bsk_tftpd_ok, 0};

  (void)op;
  (void)data;
  (void)block;
  (void)offset;
  (void)size_bytes;
  return reply;
}


typedef enum { t_unconnected=0, t_connected, t_pending} tftp_state_t;

struct bsk_tftpd_cx_t
{
  bsk_tftpd_context_info_t info;

  bsk_tftpd_log_delegate_t  log_delegate;
  bsk_tftpd_send_delegate_t send_delegate;
  bsk_tftpd_receive_delegate_t receive_delegate;

  tftp_state_t state;

  unsigned int client_port;
  unsigned int client_address;

  unsigned int block_size;

  const char* busy_msg;

  char filename[1024];
  char optmp[1024];
  char optval[1024];
  unsigned long tsize;

  unsigned int ack_buff[1024];
  tftppkt_t* ackmsg;
  unsigned int ack_len;

  int force_blocksize;
};

#define MAX_CX  8

static struct bsk_tftpd_cx_t cxdb[MAX_CX];

static void validate_delegates(bsk_tftpd_context_id_t id)
{
  bsk_tftpd_cx_t& cx=cxdb[id];

  //static struct bsk_tftpd_cx_t cx={fake_log_delegate, fake_send_delegate, fake_receive_delegate};
  if (!cx.receive_delegate)
    cx.receive_delegate=fake_receive_delegate;
  if (!cx.log_delegate)
    cx.log_delegate=fake_log_delegate;
  if (!cx.send_delegate)
    cx.send_delegate=fake_send_delegate;
}

const bsk_tftpd_context_info_t* bsk_tftpd_context_info(bsk_tftpd_context_id_t id)
{
    return &cxdb[id].info;
}

bsk_tftpd_context_id_t bsk_tftpd_context_initialize(bsk_tftpd_context_id_t id, const char* name, unsigned int port, bsk_tftp_cookie_t cookie)
{
    static int allocated_cx;
    if (id<0)
    {
        id=allocated_cx;
        ++allocated_cx;
    }
    cxdb[id].info.name=name;
    cxdb[id].info.port=port;
    cxdb[id].info.cookie=cookie;
    return id;
}

void bsk_tftpd_set_receive_delegate(bsk_tftpd_context_id_t id, bsk_tftpd_receive_delegate_t receive_delegate)
{
  cxdb[id].receive_delegate=receive_delegate ? receive_delegate : fake_receive_delegate;
}

void bsk_tftpd_set_send_delegate(bsk_tftpd_context_id_t id, bsk_tftpd_send_delegate_t send_delegate)
{
  cxdb[id].send_delegate=send_delegate ? send_delegate : fake_send_delegate;
}

void bsk_tftpd_set_log_delegate(bsk_tftpd_context_id_t id, bsk_tftpd_log_delegate_t log_delegate)
{
  cxdb[id].log_delegate=log_delegate ? log_delegate : fake_log_delegate;
}

#define LOG(msg_, ...) cx.log_delegate(id, msg_, ##__VA_ARGS__)
//#define LOG_DBG(msg_) /*cx.log_delegate msg_*/ ((void)0)
//#define LOG_DBG(msg_) cx.log_delegate msg_
#define LOG_DBG(msg_, ...) cx.log_delegate(id, msg_, ##__VA_ARGS__)


static void nakmsg(bsk_tftpd_context_id_t id, unsigned short code, const char* msg, unsigned int adr=0, unsigned int port=0)
{
  const struct errmsg {
    int e_code;
    const char *e_msg;
  } errmsgs[] = {
    { EUNDEF,   "Undefined error code" },
    { ENOTFOUND, "File not found" },
    { EACCESS,  "Access violation" },
    { ENOSPACE, "Disk full or allocation exceeded" },
    { EBADOP,   "Illegal TFTP operation" },
    { EBADID,   "Unknown transfer ID" },
    { EEXISTS,  "File already exists" },
    { ENOUSER,  "No such user" },
    { ECANCELLED, "Cancelled by administrator" },
    { -1,       0 }
 };

  if (msg==0)
  {
    int i;
    int n=0;
    for(i=0; /*i<*/errmsgs[i].e_code!=-1; ++i)
        if (errmsgs[i].e_code==code)
        {
        n=i;
        break;
        }
    msg=errmsgs[n].e_msg;
  }

  bsk_tftpd_cx_t& cx=cxdb[id];

  tftp_htons(cx.ackmsg->header.th_opcode, TFTP_ERROR);
  tftp_htons(cx.ackmsg->header.th_u.tu_code, code);
  strcpy((char*)cx.ackmsg->th_data, msg);
  if (adr==0)
      adr=cx.client_address;
  if (port==0)
      port=cx.client_port;

  cx.send_delegate(id, adr, port, cx.ackmsg, sizeof(cx.ackmsg->header)+2+strlen(msg));
  LOG("%d: %u %s (%x:%d)", id, tftp_ntos(cx.ackmsg->header.th_u.tu_code), cx.ackmsg->th_data, adr, port);
}

static void nak(bsk_tftpd_context_id_t id, unsigned short code)
{
    nakmsg(id, code, 0);
}

static int extract_str(char* dst, unsigned int dsize, char** src, unsigned int* ssize, const char* errmsg)
{
  bsk_tftpd_context_id_t id=0;

  bsk_tftpd_cx_t& cx=cxdb[id];

  char* const p=*src;
  unsigned int s=*ssize;
  unsigned int i;
  for(i=0; (i<s) && (p[i]!=0); ++i)
  {
    if (i>=(dsize-1))
    {
      LOG("TFTP: %s too long", errmsg);
      return -1;
    }
    if (p[i]>='A' && p[i]<='Z')
      dst[i]=p[i]; //disable case conversion!!!! +('a'-'A');
    else
      dst[i]=p[i];
    dst[i+1]=0;
  }
  if (i==0)
    return -1;

  s=s-i;
  if (s>1 && p[i]==0)
  {
    --s;
    ++i;
  }
  *ssize=s;
  *src=&p[i];
  return 0;
}

static void tftp_itoa(int v, char* const dst)
{
  int i;
  char str[16];
  for(i=0; v; v/=10, ++i)
  {
    int tmp=v/10;
    tmp=v-tmp*10;
    str[i]='0'+tmp;
  }
  str[i]=0;
  --i;
  if (i<0)
      i=0;
  int j;
  for(j=0; i>=0; --i, ++j)
  {
      dst[j]=str[i];
  }
  dst[j]=0;
}

static int tftp_strcmp(const char* const s1, const char* const s2)
{
  int i;
  for(i=0; s1[i] && s2[i]; ++i)
  {
    char c1=s1[i];
    if (c1>='A' && c1<='Z')
      c1=c1+('a'-'A');
    char c2=s2[i];
    if (c2>='A' && c2<='Z')
      c2=c2+('a'-'A');

    if (c1!=c2)
      return -1;
  }
  return 0;
}

static unsigned int tftp_atoi(const char* s)
{
  unsigned int value=0;
  for(; *s; ++s)
  {
    value=value*10;
    value=value+*s-'0';
  }
  return value;
}

static int doConnect(bsk_tftpd_context_id_t id, unsigned int c_adr, unsigned int c_port, const tftppkt_t* const pkt, unsigned int size)
{
  char* p=(char*)pkt->header.th_u.tu_stuff;
  unsigned int s=size-sizeof(tftphdr_t);
  int opterr=0;
  int ack_counter=0;
  int has_bklsize=0;
  bsk_tftpd_receive_status_t sts;

  bsk_tftpd_cx_t& cx=cxdb[id];

  cx.block_size=512;
  cx.tsize=0;

  memset(cx.ack_buff, 0, sizeof cx.ack_buff);
  cx.ackmsg=(tftppkt_t*)cx.ack_buff;
  tftp_htons(cx.ackmsg->header.th_opcode, TFTP_OACK);

  if (extract_str(cx.filename, sizeof cx.filename, &p, &s, "filename"))
  {
    LOG("TFTP: invalid filename");
    nak(id, EBADOP);
    return 1;
  }

  /*TODO: add grant request*/
  LOG("TFTP: filename <%s>\n", cx.filename);

  if (extract_str(cx.optmp, sizeof cx.optmp, &p, &s, "mode"))
  {
    LOG("TFTP: unsupported mode");
    nak(id, EBADOP);
    return 1;
  }

  if (tftp_strcmp(cx.optmp, "octet") && tftp_strcmp(cx.optmp, "binary"))
  {
    LOG(("TFTP: unsupported mode\n"));
    nak(id, EBADOP);
    return 1;
  }

  while(extract_str(cx.optmp, sizeof cx.optmp, &p, &s, "opt")==0)
  {
    if (extract_str(cx.optval, sizeof cx.optval, &p, &s, "optval"))
    {
      LOG("TFTP: opt=%s, value=ERROR", cx.optmp);
      nak(id, EBADOP);
      return 1;
    }

    if (tftp_strcmp(cx.optmp, TFTP_OPT_BLKSIZE)==0)
    {
      unsigned int v=tftp_atoi(cx.optval);
      LOG_DBG("TFTP: opt=%s, value=%s (%u)", cx.optmp, cx.optval, v);
      if (v<TFTP_MINSEGSIZE)
        v=TFTP_MINSEGSIZE;
      if (v>TFTP_MAXSEGSIZE)
        v=TFTP_MAXSEGSIZE;
      strcpy((char*)&cx.ackmsg->header.th_u.tu_stuff[ack_counter], TFTP_OPT_BLKSIZE);
      ack_counter+=sizeof(TFTP_OPT_BLKSIZE);
      char* p=(char*)&cx.ackmsg->header.th_u.tu_stuff[ack_counter];
      tftp_itoa(v, p);
      ack_counter+=strlen((char*)&cx.ackmsg->header.th_u.tu_stuff[ack_counter]);
      has_bklsize=1;
      cx.block_size=v;
      LOG("TFTP: block size=%d (%s)", cx.block_size, p);
    }
    else if (tftp_strcmp(cx.optmp, TFTP_OPT_TSIZE)==0)
    {
        cx.tsize=tftp_atoi(cx.optval);
        LOG("TFTP: opt=%s, value=%s", cx.optmp, cx.optval);
    }
    else
    {
      LOG("TFTP: opt=%s, value=%s - unsupported option", cx.optmp, cx.optval);
      ++opterr;
    }
  }

  if (opterr)
  {
    nak(id, EBADOP);
    return 1;
  }

  sts=cx.receive_delegate(id, bsk_tftpd_open_file, cx.filename, 0, 0, cx.tsize);
  if (sts.sts)
  {
    nakmsg(id, sts.sts, sts.ack_msg); //EBADOP);
    return 1;
  }

  if (!has_bklsize && cx.force_blocksize)
  {
    strcpy((char*)&cx.ackmsg->header.th_u.tu_stuff[ack_counter], TFTP_OPT_BLKSIZE);
    ack_counter+=sizeof(TFTP_OPT_BLKSIZE);
    tftp_itoa(TFTP_MAXSEGSIZE, (char*)&cx.ackmsg->header.th_u.tu_stuff[ack_counter]);
    ack_counter+=strlen((char*)&cx.ackmsg->header.th_u.tu_stuff[ack_counter]);
  }

  if (ack_counter==0)
  {
    tftp_htons(cx.ackmsg->header.th_opcode, TFTP_ACK);
    cx.ackmsg->header.th_u.tu_block[0]=0;
    cx.ackmsg->header.th_u.tu_block[1]=0;
    ack_counter=1;
  }

  cx.send_delegate(id, c_adr, c_port, cx.ackmsg, (ack_counter+sizeof(tftphdr_t))-1);

  return 0;
}

static int doProcessData(bsk_tftpd_context_id_t id, const tftppkt_t* const pkt, unsigned int size)
{
  bsk_tftpd_receive_status_t sts;
  unsigned short block=tftp_ntos(pkt->header.th_u.tu_block);
  unsigned int s=size-sizeof(pkt->header);

  bsk_tftpd_cx_t& cx=cxdb[id];

  LOG_DBG("TFTP: receive b%u, %ubytes\n", block, s);


  int end=s<cx.block_size;

  unsigned long offset=0;
  if (block>0)
      offset=(block-1)*cx.block_size;
  sts=cx.receive_delegate(id, end ? bsk_tftpd_end : bsk_tftpd_data, pkt->th_data, block, offset, s);
  if (sts.sts && sts.sts<bsk_tftpd_dont_ack)
  {
    nakmsg(id, sts.sts, sts.ack_msg); //EBADOP);
    return -1;
  }

  tftp_htons(cx.ackmsg->header.th_opcode, TFTP_ACK);
  cx.ackmsg->header.th_u.tu_block[0]=pkt->header.th_u.tu_block[0];
  cx.ackmsg->header.th_u.tu_block[1]=pkt->header.th_u.tu_block[1];

  if (sts.sts!=bsk_tftpd_dont_ack)
  {
    cx.send_delegate(id, cx.client_address, cx.client_port, cx.ackmsg, sizeof(cx.ackmsg->header));
  }

  if (end)
      return sts.sts!=bsk_tftpd_dont_ack ? 1 : 2;
  return 0;
}

int bsk_tftpd_process_packet(bsk_tftpd_context_id_t id, unsigned int client_address, unsigned int client_port, const void* packet, unsigned int size_bytes)
{
  //static int initialized;
  unsigned short opcode;
  tftphdr_t* hdr=(tftphdr_t*)packet;
  tftppkt_t* pkt=(tftppkt_t*)packet;

  bsk_tftpd_cx_t& cx=cxdb[id];

  //if (!initialized)
  {
    cx.ackmsg=(tftppkt_t*)cx.ack_buff;
    memset(cx.ack_buff, 0, sizeof cx.ack_buff);
    //initialized=1;
  }

  LOG_DBG("TFTP: process packet: %u %x %u", size_bytes, client_address, client_port);

  if (size_bytes<sizeof(tftphdr_t))
    return -1;

  validate_delegates(id);

  opcode=tftp_ntos(hdr->th_opcode);

  if (opcode==TFTP_ERROR)
  {
    unsigned short code=tftp_ntos(pkt->header.th_u.tu_code);
    char* p=(char*)pkt->th_data;
    unsigned int s=size_bytes-sizeof(pkt->header);
    if (!extract_str(cx.optmp, sizeof cx.optmp, &p, &s,  0))
      LOG("TFTP: error from peer: %u  %s\n", code, cx.optmp);
    else
      LOG("TFTP: error from peer: %u  (unknown)\n", code);
    if (cx.state==t_connected)
    {
      cx.receive_delegate(id, bsk_tftpd_error, 0, 0, 0, 0);
      cx.state=t_unconnected;
    }
    return -1;
  }
  if (opcode==TFTP_ACK)
  {
    LOG(("TFTP: unexpected ACK from peer\n"));
    return -1;
  }


  if (opcode==TFTP_RRQ)
  {
      char* p=(char*)pkt->header.th_u.tu_stuff;
#if 0
      unsigned int s=size_bytes-sizeof(tftphdr_t);
      if (extract_str(cx.filename, sizeof cx.filename, &p, &s, "filename"))
      {
        LOG("TFTP: invalid filename");
        nak(id, EBADOP);
        return 0;
      }
#endif
      bsk_tftpd_receive_status_t sts=cx.receive_delegate(id, bsk_tftpd_open_read_not_supported, p, 0, 0, 0);
      nakmsg(id, sts.sts, sts.ack_msg, client_address, client_port);
      return -100;
  }

  if (cx.busy_msg)
  {
      nakmsg(id, EACCESS, cx.busy_msg, client_address, client_port);
      return -10;
  }

  switch(cx.state)
  {
  case t_unconnected:
    cx.client_address=client_address;
    cx.client_port=client_port;
    if (opcode!=TFTP_WRQ)
    {
      LOG("TFTP: Unexpected request %d from peer", opcode);
      nak(id, EBADOP);
      return -2;
    }    
    if (!doConnect(id, client_address, client_port, pkt, size_bytes))
    {
      LOG("TFTP: connected: %x %u", client_address, client_port);
      cx.state=t_connected;
    }
    break;

  case t_connected:
  case t_pending:
    if (opcode==TFTP_WRQ)
    {
        if (cx.state==t_pending)
        {
            nakmsg(id, EACCESS, "pending operation");
            return 2;
        }
    }
    if ((cx.client_address!=client_address) || (cx.client_port!=client_port))
    {
        LOG("TFTP: ERROR - already connected!");
        nak(id, EBADID);
        return -5;
    }
    if (opcode==TFTP_DATA)
    {
      if (cx.state==t_pending)
      {
          return 2;
      }
      int res=doProcessData(id, pkt, size_bytes);
      if (res<0)
      {
        LOG("TFTP: ERROR - disconnect!\n");
        cx.state=t_unconnected;
        return -3;
      }
      if (res>0)
      {
          if (res==2)
          {
              cx.state=t_pending;
              LOG("TFTP: completed - pending ack\n");
          }
          else
          {
              LOG("TFTP: completed (%d blocks)", 0);
            cx.state=t_unconnected;
          }
        return res;
      }
    }
    else
    {
      LOG("TFTP: disconnect - invalid opcode: %u\n", opcode);
      nak(id, EBADOP);
      cx.state=t_unconnected;
      return -4;
    }
    break;
  }

  return 0;
}



int bsk_tftpd_send_ack(bsk_tftpd_context_id_t id, bsk_tftpd_status_t /*code*/, const char* /*msg*/)
{
    bsk_tftpd_cx_t& cx=cxdb[id];

    LOG("ACK!");
    tftp_htons(cx.ackmsg->header.th_opcode, TFTP_ACK);
    cx.ackmsg->header.th_u.tu_block[0]=0;
    cx.ackmsg->header.th_u.tu_block[1]=0;
    cx.send_delegate(id, cx.client_address, cx.client_port, cx.ackmsg, sizeof(cx.ackmsg->header));
    cx.state=t_unconnected;

    return 0;
}

int bsk_tftpd_set_busy(bsk_tftpd_context_id_t id, const char* msg)    //msg=0 means NO BUSY
{
    bsk_tftpd_cx_t& cx=cxdb[id];

    cx.busy_msg=msg;

    return 0;
}


//Client I/F