/*

                                   _JNJ`                        
                               .JNMH`                           
                              JMMF`       `;.                   
                            .NMM)          `MN.                 
                            MMM)            (MML                
                           (MMM`             MMML               
                     M  I  D  N  I  G  H  T    C  o  D  E       
                           (NMMF             MHNH               
                            NMML            .MMM                
                             NMML          .NMH                 
                              4MMNL       .#F                   
                               `4HNNL_   `                      
                                   `"""`                        

                          Copyright (C) 2004-2007
       "Ian (Larry) Latter" <ian dot latter at midnightcode dot org>

    Midnight Code is a registered trademark of Ian Latter.

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, and mirrored at the Midnight Code web
    site; as at version 2 of the License only.
  
    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
    General Public License for more details.
  
    You should have received a copy of the GNU General Public License
    (version 2) along with this program; if not, write to the Free
    Software Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
    02111-1307 USA, or see http://midnightcode.org/gplv2.txt

*/

/* midnightcode-tftp.c :: Midnight Code TFTP library			*/

/* Includes software to support RFCs :: 783, 1350			*/


#include <stdio.h>
#include <string.h>
#include <stdarg.h>
#include <stdlib.h>
#include <unistd.h>
#include <netdb.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <arpa/tftp.h>



#include "preferences.h"
#include "midnightcode-list.h"
#include "midnightcode-config.h"
#include "midnightcode-socket.h"
#include "midnightcode-tftp.h"



int
midnightcode_tftp_get(char ** buffer, int * buffer_size,
                        char * target_host, int target_port,
                        char * file_path) {
  struct midnightcode_socket_element_t * socket;
  struct tftphdr * tftp_send;
  struct tftphdr * tftp_receive;
  char * tftp_receive_p;
  char * tftp_send_p;
  int source_port;
  int bytes_read;
  int tftp_size;
  int tftp_len;
  int timeout;
  int timeout_count;
  int timeout_retries;
  int successful_block;
  struct sockaddr_in sockaddr;
  int sockaddr_len;
  struct hostent *host_entry;



  if(!target_host)
    return -1;
  if(target_port < 0 || target_port > 65535)
    return -1;
  if(!file_path)
    return -1;
  if(strlen(file_path) > (SEGSIZE - 16))
    return -1;


  *buffer_size = 0;

  tftp_size = sizeof(struct tftphdr) + SEGSIZE;
  tftp_send = malloc(tftp_size);
  if(!tftp_send)
    return -1;

  if(target_port == 0)
    target_port = MIDNIGHTCODE_TFTP_CONTROL_PORT;

  source_port = MIDNIGHTCODE_TFTP_DATA_PORT;
  while(source_port < 65535) {
    if(midnightcode_socket_open(&socket,
      MIDNIGHTCODE_SOCKET_CONTROL_RAW, 0, 1,
      MIDNIGHTCODE_SOCKET_DIR_LISTEN, MIDNIGHTCODE_SOCKET_PROTOCOL_UDP,
      "0.0.0.0", source_port) < 0) {
      source_port++;
      continue;
    }
    break;
  }
  if(source_port == 65535) {
    free(tftp_send);
    return -1;
  }

  socket->port_remote = target_port;
  host_entry = NULL;
  if((host_entry = gethostbyname(target_host)) == NULL) {
    free(tftp_send);
    return -1;
  }
  memcpy(&socket->ipv4_remote, host_entry->h_addr, sizeof(struct in_addr));
  snprintf(socket->ipv4_remote_s, 32, "%s", inet_ntoa(socket->ipv4_remote));

  memset(&sockaddr, 0, sizeof(struct sockaddr_in));
  sockaddr.sin_family = AF_INET;
  sockaddr.sin_addr.s_addr = socket->ipv4_remote.s_addr;
  sockaddr.sin_port = htons(socket->port_remote);
  sockaddr_len = sizeof(struct sockaddr_in);

  /* TODO: Implement RFCs 2347 and 2348 for "blksize" support */

  memset(tftp_send, 0, tftp_size);
  tftp_send->th_opcode = htons(RRQ);
  memcpy(&(tftp_send->th_stuff), file_path, strlen(file_path));
  tftp_send_p = (char *)(tftp_send->th_stuff);
  tftp_send_p += strlen(file_path) + 1;
  memcpy(tftp_send_p, "OCTET", strlen("OCTET"));
  tftp_len = 3 + strlen(file_path) + 1 + strlen("OCTET") + 1;

  tftp_send_p = (char *)(tftp_send);
  if(sendto(socket->socket, tftp_send_p, tftp_len, 0,
    (struct sockaddr *)&sockaddr, sockaddr_len) < 0) {
    midnightcode_socket_close(socket);
    free(tftp_send);
    return -1;
  }

  timeout = 0;
  timeout_retries = 0;
  timeout_count = MIDNIGHTCODE_TFTP_TIMEOUT *
    (1000000 / MIDNIGHTCODE_TFTP_POLL_TIME);

  successful_block = 0;

  while(1) {

    if(timeout >= timeout_count) {

      timeout = 0;
      timeout_retries++;
      if(timeout_retries > MIDNIGHTCODE_TFTP_RETRIES) {
        midnightcode_socket_close(socket);
        free(tftp_send);
        return -1;
      }

      if(successful_block == 0) {

        memset(tftp_send, 0, tftp_size);
        tftp_send->th_opcode = htons(RRQ);
        memcpy(&(tftp_send->th_stuff), file_path, strlen(file_path));
        tftp_send_p = (char *)(tftp_send->th_stuff);
        tftp_send_p += strlen(file_path) + 1;
        memcpy(tftp_send_p, "OCTET", strlen("OCTET"));
        tftp_len = 3 + strlen(file_path) + 1 + strlen("OCTET") + 1;

        tftp_send_p = (char *)(tftp_send);
        if(sendto(socket->socket, tftp_send_p, tftp_len, 0,
          (struct sockaddr *)&sockaddr, sockaddr_len) < 0) {
          midnightcode_socket_close(socket);
          free(tftp_send);
          return -1;
        }

      } else {

        memset(tftp_send, 0, tftp_size);
        tftp_send->th_opcode = htons(ACK);
        tftp_send->th_block = htons(successful_block);
        tftp_len = 4;

        tftp_send_p = (char *)(tftp_send);
        if(sendto(socket->socket, tftp_send_p, tftp_len, 0,
          (struct sockaddr *)&sockaddr, sockaddr_len) < 0) {
          midnightcode_socket_close(socket);
          free(tftp_send);
          return -1;
        }

      }
      usleep(MIDNIGHTCODE_TFTP_POLL_TIME);

    }

    tftp_size = sizeof(struct tftphdr) + SEGSIZE;
    tftp_receive_p = NULL;
    bytes_read = 0;
    if(midnightcode_socket_read_block(socket,
      &tftp_receive_p, &bytes_read,
      1, MIDNIGHTCODE_SOCKET_IO_NONBLOCKING) <= 0) {
      timeout++;
      usleep(MIDNIGHTCODE_TFTP_POLL_TIME);
      continue;
    }

    if(bytes_read > tftp_size)
      bytes_read = tftp_size;
    tftp_receive = (struct tftphdr *)tftp_receive_p;

    if(successful_block == 0)
      sockaddr.sin_port = htons(socket->port_remote);

    switch(ntohs(tftp_receive->th_opcode)) {

      case DATA:
        if(ntohs(tftp_receive->th_block) == (successful_block + 1)) {
          tftp_len = bytes_read - 4;
          fflush(stdout);
          *buffer = realloc(*buffer, (*buffer_size) + tftp_len);
          if(!(*buffer)) {
            midnightcode_socket_close(socket);
            free(tftp_send);
            free(tftp_receive);
            return -1;
          }
          memcpy(((*buffer) + *buffer_size), &(tftp_receive->th_data),
            tftp_len);
          *buffer_size += tftp_len;
          successful_block++;
        }
        free(tftp_receive);
        tftp_receive = NULL;

        memset(tftp_send, 0, tftp_size);
        tftp_send->th_opcode = htons(ACK);
        tftp_send->th_block = htons(successful_block);
        tftp_len = 4;

        tftp_send_p = (char *)(tftp_send);
        if(sendto(socket->socket, tftp_send_p, tftp_len, 0,
          (struct sockaddr *)&sockaddr, sockaddr_len) < 0) {
          midnightcode_socket_close(socket);
          free(tftp_send);
          return -1;
        }

        if((bytes_read - 4) < SEGSIZE) {
          midnightcode_socket_close(socket);
          free(tftp_send);
          return *buffer_size;
        }
        break;

      case ERROR:
        midnightcode_socket_close(socket);
        free(tftp_send);
        free(tftp_receive);
        return -1;
        break;

      default:
        midnightcode_socket_close(socket);
        free(tftp_send);
        free(tftp_receive);
        return -1;
        break;

    }
  }


  return *buffer_size;
}


