/*
 * $Id: netcat.c,v 1.20 2001/12/29 16:40:15 coelho Exp $
 *
 * (c) Fabien Coelho <fabien@coelho.net> 2001
 *
 * LICENSE: GNU-GPL (see http://www.gnu.org/)
 * NO WARRANTY WHATSOEVER.
 * PLEASE USE THIS PROGRAM UNDER PROPER SUPERVISION BY A NETWORK ADMINISTRATOR.
 *
 * You can do whatever you want with this source, provided that this 
 * copyright is preserved and aknowledged. If you modify this source,
 * the modifications must be properly documented and signed. 
 *
 * my 'netcat' implementation. 
 * no other documentation than this source code.
 * see 'command -h' for help about options.
 * see 'man nc' for the netcat tcp/udp server/client program.
 */

/**************************************************************** C INCLUDES */

#include <stdio.h>  /* fprintf FILE */
#include <string.h> /* strchr strcmp */
#include <stdlib.h> /* atoi */

/****************************************************** UNIX SYSTEM INCLUDES */

#include <sys/types.h>
#include <sys/time.h>
#include <sys/socket.h>

#include <unistd.h> /* read write select */

#include <netinet/in.h>
#include <netdb.h>
#include <signal.h>
#include <errno.h>

/**************************************************** CONVENIENT DEFINITIONS */

typedef enum { false, true } boolean;

typedef char * string;

#define DEFAULT_BUFFER_SIZE (65536)

#define DEFAULT_PORT "80"

/************************************************************ GLOBAL OPTIONS */

/* verbose about what is being done. */
static boolean verbose = false;

/* debug messages on. */
static boolean debug = false;

/* whether to reconnect on loops. */
static boolean reconnect = false;

/* whether to only display printable characters. */
static boolean printable = false;

/* seconds to wait between files. */
static int seconds = 0;

/* number of seconds without activity from server before timeout. */
static int timeout = -1;

/* how many loops over the files. */
static int loop = 1;

/* internal buffer size. */
static int buffer_size = DEFAULT_BUFFER_SIZE;

/********************************************************************* UTILS */

static void set_fd(int fd, fd_set * pfd, int * pn)
{
  FD_SET(fd, pfd);
  if (fd>*pn) *pn = fd;
}

/* the buffer. */
static char * buffer = NULL;

/* copy in to out if needed. return whether empty (=> close?)
 */
static boolean move(int in, int out, boolean printed)
{
  ssize_t rsize, wsize;

  if (debug) fprintf(stderr, "reading %d\n", in);
  
  if ((rsize = read(in, buffer, buffer_size))==-1)
  {
    perror("read()");
    exit(11);
  }
  
  if (debug) fprintf(stderr, " - %d bytes read\n", rsize);
  
  if (rsize==0)
    return false;
  
  if (printed) {
    int i;
    for (i=0; i<rsize; i++)
      if ((buffer[i]<32 && buffer[i]!=10) || (buffer[i] > 126))
	buffer[i] = '.';  /* keep 10 + 32..126 */
  }

  if (debug) fprintf(stderr, "writing %d\n", out);
  
  if ((wsize = write(out, buffer, rsize))==-1)
  {
    perror("write()");
    exit(12);
  }
  
  if (wsize!=rsize)
  {
    fprintf(stderr, "wsize!=rsize... what should I do?\n");
    exit(16);
  }

  if (debug) fprintf(stderr, " - %d bytes written\n", wsize);

  return true;
}

/* help message */
static void usage(string program, int exitcode)
{
  fprintf(stderr, 
    "Usage: %s [-?hdprv] [-l n] [-s s] [-t t] [-b b] host[:port] [files...]\n"
	  "\tsend files to network $Revision: 1.20 $\n"
	  "\t-h or -?: this help\n"
	  "\t-v: verbose (default false)\n"
	  "\t-d: debug (default false)\n"
	  "\t-r: reconnect on loops (default false)\n"
	  "\t-p: print printable characters only (default false)\n"
	  "\t-l n: loop on files n times (for ever 0, default 1)\n"
	  "\t-s s: sleep s seconds between files (default 0)\n"
	  "\t-t t: timeout in seconds without events (default -1 for none)\n"
	  "\t-b b: internal buffer size in bytes (default %d)\n"
	  "\thost: destination host\n"
	  "\tport: destination tcp port (default %s)\n"
	  "\tfiles...: files to send, '-' for stdin (default stdin)\n",
	  program, DEFAULT_BUFFER_SIZE, DEFAULT_PORT);
  exit(exitcode);
}

/**************************************************************** DO THE JOB */

int main(int argc, string argv[])
{
  int c, index, number=1;
  int out = fileno(stdout), dst = -1, nin = -1;
  string dst_host, dst_port;
  struct sockaddr_in dst_addr;
  FILE * in = NULL;
  boolean stdin_used = false, in_r, dst_w, dst_r, out_w;

  /* option management */
  while ((c=getopt(argc, argv, "?hdb:l:rps:t:v"))!=EOF)
  {
    switch (c)
    {
    case 'v':
      verbose = true;
      break;
    case 'd':
      debug = true;
      break;
    case 'r':
      reconnect = true;
      break;
    case 'p':
      printable = true;
      break;
    case 'l':
      loop = atoi(optarg);
      break;
    case 's':
      seconds = atoi(optarg);
      break;
    case 't':
      timeout = atoi(optarg);
      break;
    case 'b':
      buffer_size = atoi(optarg);
      break;
    case 'h':
    case '?':
    default: 
      usage(argv[0], 0);
    }
  }

  if (optind==argc) 
    usage(argv[0], 1);

  if (buffer_size<=0)
    usage(argv[0], 6);

  buffer = (char*) malloc(buffer_size);
  if (!buffer)
  {
    perror("malloc()");
    exit(15);
  }

  dst_host = argv[optind++];
  dst_port = strchr(dst_host, ':');
  if (dst_port) *dst_port++ = '\0';
  else dst_port = DEFAULT_PORT;

  /* get host ip */
  {
    struct hostent * he = gethostbyname(dst_host);
    if (!he) 
    {
      herror("gethostbyname()");
      exit(2);
    }

    dst_addr.sin_family = AF_INET;
    dst_addr.sin_port = htons(atoi(dst_port));
    dst_addr.sin_addr.s_addr = *((unsigned int *)he->h_addr_list[0]);
  }

  if (optind==argc)
  {
    loop = 1, in = stdin, stdin_used = true;
    if (verbose)
      fprintf(stderr, "# default '-' (stdin)\n");
  }

  index = optind;

  in_r = false;
  dst_w = false;
  dst_r = false;
  out_w = false;

  if (debug)
    fprintf(stderr, 
	    "debug=%d verbose=%d loop=%d index=%d optind=%d argc=%d\n",
	    debug, verbose, loop, index, optind, argc);

  /* main loop */
  while (true)
  {
    fd_set toread, towrite;
    struct timeval tout;
    int n, code;

    /* file to send. */
    if (!in)
    {
      if (index==argc && (loop==0 || number<loop))
      {
	if (seconds)
	{
	  if (verbose) fprintf(stderr, "# sleep %d\n", seconds);
	  sleep(seconds);
	}

	index = optind;
	number++;

	if (reconnect && dst!=-1)
	{
	  if (verbose) fprintf(stderr, "# closing...\n");

	  if (shutdown(dst, 2))
	  {
	    perror("shutdown()");
	    exit(13);
	  }

	  if (close(dst))
	  {
	    perror("close()");
	    exit(14);
	  }

	  dst = -1;
	  dst_r = false;
	  dst_w = false;
	}
      }
    }
      
    /* connexion */
    if (dst==-1)
    {
      if (verbose) 
	  fprintf(stderr, "# connecting to %s:%s\n",
		  dst_host, dst_port);
      
      if ((dst = socket(AF_INET, SOCK_STREAM,  0)) == -1)
      {
	perror("socket()");
	exit (3);
      }
	
      if (connect(dst, (struct sockaddr *) &dst_addr, 
		  sizeof(struct sockaddr)))
      {
	perror("connect()");
	exit(4);
      }
      
      dst_r = false;
      dst_w = false;
    }
    
    if (index==optind && !in && optind!=argc && verbose)
      fprintf(stderr, "# loop %d/%d\n", number, loop);

    if (!in && index<argc)
    {
      if (verbose)
	fprintf(stderr, "# file '%s' arg %d\n", argv[index], index);
      
      if (!strcmp(argv[index], "-"))
      {
	if (!stdin_used) 
	  in = stdin, stdin_used = true;
	else
	  if (verbose)
	    fprintf(stderr, "# stdin already used, skipping\n");
      }
      else
      {
	in = fopen(argv[index], "r");
	if (!in)
	{
	  perror("fopen()");
	  exit(5);
	}
      }
      
      index++;
      in_r = false;
    }

    nin = in? fileno(in): -1;

    /* wait for anything.
     * if ready to read, check if ready to write.
     */
    n = 0;
    FD_ZERO(&toread);
    FD_ZERO(&towrite);
    if (!dst_r)               set_fd(dst, &toread,  &n);
    if (dst_r && !out_w)      set_fd(out, &towrite, &n);
    if (in && !in_r)          set_fd(nin, &toread,  &n);
    if (in && in_r && !dst_w) set_fd(dst, &towrite, &n);
    n++;

    tout.tv_sec = timeout, tout.tv_usec = 0;

    if (debug)
      fprintf(stderr, 
	      "n=%d out=%d dst=%d nin=%d\tRin=%d Rdst=%d Wout=%d Wdst=%d\n",
	      n, out, dst, nin, in? in_r: 0, dst_r, out_w, dst_w);

    code = select(n, &toread, &towrite, NULL, 
		  (timeout!=-1 && !in)? &tout: NULL);

    if (code==-1 && errno!=EINTR)
    {
      perror("select()");
      exit(7);
    }

    if (code==-1 && errno==EINTR)
    {
      /* just an interrupt which made the select return: ignore */
      FD_ZERO(&toread);
      FD_ZERO(&towrite);
    }

    if (in && FD_ISSET(nin, &toread)) in_r  = true;
    if (FD_ISSET(dst, &toread))       dst_r = true;
    if (FD_ISSET(dst, &towrite))      dst_w = true;
    if (FD_ISSET(out, &towrite))      out_w = true;

    if (debug) 
      fprintf(stderr, "select()=%d\tRin=%d Rdst=%d Wout=%d Wdst=%d\n",
	      code, in? in_r: 0, dst_r, out_w, dst_w);
    
    /* back FIRST!!! 
     * funny bug here: it is not because select told you that two
     * streams are ready to be written and read that you can do it safely!
     * the first write may take resources which will make the second
     * read/write hang, especially when dealing with network.
     * hence unexpected blocking occurs, even after careful selects. 
     * So read before write?
     * or is it a bug in the linux loopback interface that I use for tests?
     * or elsewhere?
     */
    if ((dst_r && out_w && (dst_r=false, out_w=false, 
			    !move(dst, out, printable))) || 
	(!code && !in))
    {
      if (verbose) fprintf(stderr, "# closing connexion\n");

      if (shutdown(dst, 2))
	perror("shutdown()");
      
      if (close(dst))
	perror("close()");
      
      if (in && fclose(in))
	perror("fclose()");
      in=NULL;
      
      exit(9);
    }

    /* go */
    if (in && in_r && dst_w)
    {
      if (!move(nin, dst, false))
      {
	if (fclose(in))
	{
	  perror("fclose()");
	  exit(8);
	}
	in=NULL; 
      }
      in_r = false;
      dst_w = false;
    }
  }

  return 10; /* never reached. */
}