/* Copyright 2001  Mark Pulford <mark@kyne.com.au>
 * This file is subject to the terms and conditions of the GNU General Public
 * License. Read the file COPYING found in this archive for details, or
 * visit http://www.gnu.org/copyleft/gpl.html
 */

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/socket.h>
#include <fcntl.h>
#include <netdb.h>
#include <errno.h>
#include <unistd.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>

#ifdef HAVE_GETOPT_H
#include <getopt.h>
#endif

/* Configuration vars */
char *output_name = "-";
int output_fd;
int block_size = 512;
int verbose = 0;
int tcp_nodelay = 0;
#ifdef TCP_CORK
int tcp_cork = 0;
#endif

int total_written = 0;

int make_connect(const char *host, int port);
int wait_connect(int port);
int open_stream(const char *name, int write);
int resolv_addr(struct in_addr *in, const char *addr);
int output(int rfd, char *rname);
int safe_write(int fd, const char *buf, int size);
void init(int argc, char **argv);
char next_arg(int argc, char **argv);
void set_tcp_opt(int fd, int opt);
int size2num(int *size, const char *s);
void usage();
void version();
void warn(const char *msg, ...);
#ifndef HAVE_HSTRERROR
#define hstrerror(a) strherror(a)
char *strherror(int e);
#endif
#ifndef HAVE_INET_ATON
int inet_aton(const char *cp, struct in_addr *inp);
#endif

int main(int argc, char **argv)
{
	int rfd;
	int i;
	int err = 0;

	init(argc, argv);

	if(optind < argc) {
		for(i=optind; i<argc; i++) {
			rfd = open_stream(argv[i],0);
			if(rfd < 0 || !output(rfd, argv[i]))
				err++;
		}
	} else {
		rfd = open_stream("-", 0);
		if(!output(rfd, "stdin"))
			err++;
	}

	if(verbose)
		warn("%s:\t%d bytes written", output_name, total_written);

	return err?1:0;
}

/* Init output_fd and block size or die trying */
void init(int argc, char **argv)
{
	int ch;
	
	ch = next_arg(argc, argv);
	while(ch != -1) {
		switch(ch) {
		case ':': case '?':
			usage();
			exit(-1);
		case 'h':
			usage();
			exit(0);
		case 'V':
			version();
			exit(0);
		case 'v':
			verbose++;
			break;
		case 'b':
			if(!size2num(&block_size, optarg)) {
				warn("bad number given for --block-size");
				exit(-1);
			}
			break;
		case 'o':
			output_name = optarg;
			break;
		case 'n':
			tcp_nodelay = 1;
			break;
#ifdef TCP_CORK
		case 'c':
			tcp_cork = 1;
			break;
#endif
		default:
			abort();	/* BUG: Unimplemented option */
		}
		ch = next_arg(argc, argv);
	}
#ifdef TCP_CORK
	if(tcp_nodelay && tcp_cork) {
		warn("--cork and --no-delay cannot be used together");
		exit(-1);
	}
#endif
	output_fd = open_stream(output_name, 1);
	if(output_fd < 0)
		exit(-1);
}

char next_arg(int argc, char **argv)
{
	char *optstr = "b:o:Vvhn"
#ifdef TCP_CORK
	"c"
#endif
	;
#ifdef HAVE_GETOPT_H
	struct option optslong[] = {
		{"block-size", 1, NULL, 'b'},
		{"output", 1, NULL, 'o'},
		{"version", 0, NULL, 'V'},
		{"verbose", 0, NULL, 'v'},
		{"help", 0, NULL, 'h'},
		{"no-delay", 0, NULL, 'n'},
#ifdef TCP_CORK
		{"cork", 0, NULL, 'c'},
#endif
		{NULL, 0, NULL, 0}
	};

	return getopt_long(argc, argv, optstr, optslong, NULL);
#else
	return getopt(argc, argv, optstr);
#endif
}

void usage()
{
	fprintf(stderr, "Usage: ncat [OPTION...] [STREAM...]\n");
	fprintf(stderr, "  -o  --output=STREAM  Output to STREAM instead of standard output\n");
	fprintf(stderr, "  -b  --block-size=N   Change block size for reading and writing\n");
	fprintf(stderr, "  -v  --verbose        Show transfer statistics\n");
	fprintf(stderr, "  -n  --no-delay       Disable tcp buffering of small packets\n");
#ifdef TCP_CORK
	fprintf(stderr, "  -c  --cork           Buffer partial frames\n");
#endif
	fprintf(stderr, "  -V  --version        Display version\n");
	fprintf(stderr, "  -h  --help           Display help\n");
}

void version()
{
	fprintf(stderr, "ncat v%s\n", VERSION);
}

/* Return:	1	success (rfd closed)
 *              0	failure (rfd closed)
 * output() exits on fatal error */
int output(int rfd, char *rname)
{
	int total_read = 0;
	int rsize, wsize;
	static char *buf = NULL;

	if(!buf) {
		buf = malloc(block_size);
		if(!buf) {
			warn("%s: Out of memory", rname);
			exit(-1);
		}
	}

	while((rsize = read(rfd, buf, block_size))) {
		if(rsize < 0) {
			if(errno == EINTR) {
				continue;
			} else {
				warn("%s: %s", rname, strerror(errno));
				close(rfd);
				return 0;
			}
		}
		total_read += rsize;
		wsize = safe_write(output_fd, buf, rsize);
		total_written += wsize;
		if(wsize != rsize) {
			warn("write during %s: %s", rname, strerror(errno));
			exit(-1);
		}
	}

	if(verbose)
		warn("%s:\t%d bytes read", rname, total_read);

	close(rfd);
	return 1;
}

/* Returns bytes written, if != size an error occured */
int safe_write(int fd, const char *buf, int size)
{
	int i = 0;
	int ret;

	while (i < size) {
		ret = write(fd, buf+i, size-i);
		if(ret < 0) {
			if(EINTR == errno)
				continue;
			else
				return i;
		}
		i += ret;
	}

	return i;
}

/* returns -1 on error */
int open_stream(const char *name, int write)
{
	char *port;
	char *host;
	int fd;

	host = strdup(name);
	if(!host) {
		warn("%s: Out of memory", name);
		return -1;
	}

	port = strrchr(host, ':');

	/* If the stream ends with : force a filename (provided the
	 * filename is at least 1 character long) */
	if(port && !port[1] && strlen(host)>1) {
		*port = 0;
		port = NULL;
	}

	if(port == NULL) {
		/* open file */
		if(!strcmp(host, "-")) {
			if(write)
				fd = fileno(stdout);
			else
				fd = fileno(stdin);
		} else {
			if(write)
				fd = open(host, O_WRONLY|O_CREAT, 0666);
			else
				fd = open(host, O_RDONLY);
		}
		if(fd < 0)
			warn("%s: %s", host, strerror(errno));
	} else {
		/* TCP stream requested */
		if(port == host) {
			/* IF the stream is ':', then atoi("") returns 0 */
			fd = wait_connect(atoi(port+1));
		} else {
			*port = 0;
			port++;
			fd = make_connect(host, atoi(port));
		}
		if(fd < 0)
			return -1;
		if(write) {
#ifdef TCP_CORK
			/* If possible only send full frames */
			if(tcp_cork)
				set_tcp_opt(fd, TCP_CORK);
			else
#endif
			if(tcp_nodelay)
				set_tcp_opt(fd, TCP_NODELAY);
		}
		/* Shutdown unused side of the connection */
		if(shutdown(fd, !write) < 0) {
			warn("%s: %s", name, strerror(errno));
			return -1;
		}
	}

	free(host);

	return fd;
}

void set_tcp_opt(int fd, int opt)
{
	int on = 1;
	
	if(setsockopt(fd, IPPROTO_TCP, opt, &on, sizeof(on)) < 0) {
		warn("setsockopt: %s", strerror(errno));
		exit(-1);
	}
}

int make_connect(const char *host, int port)
{
	struct sockaddr_in addr;
	int fd;

	if(!resolv_addr(&addr.sin_addr, host)) {
		warn("%s:%d: %s", host, port, hstrerror(h_errno));
		return -1;
	}

	addr.sin_family = AF_INET;
	addr.sin_port = htons(port);

	fd = socket(PF_INET, SOCK_STREAM, 0);
	if(fd < 0) {
		warn("%s:%d: %s", host, port, strerror(errno));
		return -1;
	}

	if(connect(fd, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
		warn("%s:%d: %s", host, port, strerror(errno));
		close(fd);
		return -1;
	}

	return fd;
}

int resolv_addr(struct in_addr *in, const char *addr)
{
	struct hostent *h;

	if(inet_aton(addr, in))
		return 1;

	h = gethostbyname(addr);
	if(!h)
		return 0;
	memcpy(in, h->h_addr_list[0], sizeof(*in));

	return 1;
}

int wait_connect(int port)
{
	struct sockaddr_in addr;
	int fd;
	int nfd;
	int size;
	int on = 1;

	fd = socket(PF_INET, SOCK_STREAM, 0);
	if(fd < 0) {
		warn(":%d: %s", port, strerror(errno));
		return -1;
	}

	setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on));

	addr.sin_family = AF_INET;
	addr.sin_addr.s_addr = INADDR_ANY;
	addr.sin_port = htons(port);

	if(bind(fd, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
		warn(":%d: %s", port, strerror(errno));
		close(fd);
		return -1;
	}

	if(listen(fd, 1) < 0) {
		warn(":%d: %s", port, strerror(errno));
		close(fd);
		return -1;
	}

	if(0 == port) {
		size = sizeof(addr);

		if(getsockname(fd, (struct sockaddr *)&addr, &size) < 0) {
			warn(":%d: %s", port, strerror(errno));
			close(fd);
			return -1;
		}
		warn("listening on port %d", ntohs(addr.sin_port));
	}

	size = sizeof(addr);
	nfd = accept(fd, (struct sockaddr *)&addr, &size);
	if(nfd < 0) {
		warn(":%d: %s", port, strerror(errno));
		close(fd);
		return -1;
	}
	close(fd);	/* Close server sock */

	return nfd;
}

void warn(const char *msg, ...)
{
	va_list arg;

	fprintf(stderr, "ncat: ");
	va_start(arg, msg);
	vfprintf(stderr, msg, arg);
	va_end(arg);
	fprintf(stderr, "\n");
}

/* handles b, k, m suffixes
 * returns: 0	failure
 * 	    1	success
 * Note: this doesn't handle overflow */
int size2num(int *size, const char *s)
{
	char *suff;

	*size = strtol(s, &suff, 10);
	if(suff == s)
		return 0;	/* No number found */
	if(*size < 1)
		return 0;	/* Must have positive non zero */

	if(0 == *suff)
		return 1;	/* No suffix */

	if(0 != suff[1])
		return 0;	/* suffix too long */

	switch(*suff) {
	case 'b': case 'B':
		*size *= 512;
		return 1;
	case 'k': case 'K':
		*size *= 1024;
		return 1;
	case 'm': case 'M':
		*size *= 1048576;
		return 1;
	}

	return 0;	/* Bad modifier */
}

#ifndef HAVE_INET_ATON
int inet_aton(const char *cp, struct in_addr *inp)
{
	inp->s_addr = inet_addr(cp);

	return -1 != inp->s_addr;
}
#endif

#ifndef HAVE_HSTRERROR
char *strherror(int e)
{
	switch(e) {
	case HOST_NOT_FOUND:
		return "host not found";
	case TRY_AGAIN:
		return "server failure";
	case NO_RECOVERY:
		return "non recoverable error during lookup";
	case NO_DATA:
		return "no such address";
	default:
		return "unknown error resolving address";
	}
}
#endif