/*
	 sockpipe.c
	 gcc -o sockpipe sockpipe.c
	 
	 binds to a port on localhost and forwards all information to remote host
	 and port
	 
	 by hamblin@cs.wisc.edu
*/
	 
#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <errno.h>
#include <unistd.h>
#include <signal.h>
#include <string.h>
#include <utmp.h>
#include <netdb.h>
#define LOGS "/tmp"
extern int errno;
extern int h_errno;


int debug = 0;
int local_port, remote_port;
char remote_host[80];

int start_listening(int *s, struct sockaddr_in *sa);
int wait_for_conns(int s, struct sockaddr_in *remote, int flag);
void sig_chld(int sig);


int main(int argc, char **argv)
{
	int log_flag = 0;
	int sock, n, remote_size;
	struct sockaddr_in serv_addr, remote_addr;
	struct hostent *he; 
	
	if (argc < 4)
	{
		printf("usage: %s local_port remote_host remote_port [log 1/0]\n", argv[0]);
		exit(-1);
	}
	
	local_port = atoi(argv[1]);
	remote_port = atoi(argv[3]);

	if (argc == 5)
		log_flag = atoi(argv[4]);

	/* make sure bs->requestee.host is an ip address */

	if (!isdigit(argv[2][0]))
	{
		if ( (he = gethostbyname(argv[2])) == NULL)
			return -1;
		if (inet_ntop(he->h_addrtype, he->h_addr_list[0], remote_host, 80) == NULL)
		{
			perror("inet_ntop");
			exit(-1);
		}
	}
	else
		strncpy(remote_host, argv[2], 80);

	printf("forwarding connections to localhost:%d to %s:%d\n", 
			local_port, remote_host, remote_port);


	/* build remote_addr */
	bzero((char *) &remote_addr, sizeof(remote_addr));
	remote_size = sizeof(remote_addr);
	remote_addr.sin_family = AF_INET;
	remote_addr.sin_port = htons(remote_port);
	if (inet_pton(AF_INET, remote_host, &remote_addr.sin_addr) == -1)
	{
		perror("inet_pton");
		exit(-1);
	}

	if (start_listening(&sock, &serv_addr) == -1)
	{
		fprintf(stderr, "Failure in start_listening.\n");
		exit(-1);
	}
	if (debug)
			printf("started listening\n");

	if (debug == 0)
	{
		if (fork() != 0)
			exit(0);
		if (setsid() == -1)
		{
			perror("setsid");
			exit(-1);
		}
		signal(SIGHUP, SIG_IGN);
		signal(SIGCHLD, sig_chld);
		if (fork() != 0)
			exit(0);
		chdir("/");
		umask(0);
		for (n = 0; n < 3; n ++)
			close(n);
	}

	while (1)
	{
		if (wait_for_conns(sock, &remote_addr, log_flag) == -1)
		{
			printf("error in wait_for_conn\n");
			exit(-1);
		}
	}
}

int start_listening(int *s, struct sockaddr_in *sa)
{
	int option = 1;

	if ( (*s = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)) == -1)
	{
		perror("socket");
		return -1;
	}
	if (debug)
		printf("start_listening: got socket %d\n", *s);

	/* set up addr */
	bzero((char *) sa, sizeof(*sa));
	sa->sin_family = AF_INET;
	sa->sin_port = htons(local_port);
	sa->sin_addr.s_addr = htonl(INADDR_ANY);

	/* set SO_REUSEADDR so that will always bind to the addr */
	if (setsockopt(*s, SOL_SOCKET, SO_REUSEADDR, &option, sizeof(int)) == -1)
	{
		perror("setsockopt");
		return -1;
	}

	/* bind socket to name */	
	if (bind(*s, (struct sockaddr *) sa, sizeof(*sa)) == -1)
	{
		perror("bind");
		return -1;
	}
	if (debug)
		printf("start_listening: socket bound.\n");

	return 1;
}

int wait_for_conns(int s, struct sockaddr_in *remote, int flag)
{
	int accepted_sock, forward_sock, cli_size, max_fds = s, n;
	struct sockaddr_in cli_addr;
	fd_set sockets;
	char buf[1024];
	FILE *log = 0;
	char log_name[80];
	
	bzero((char *) &cli_addr, sizeof(cli_addr));
	cli_size = sizeof(struct sockaddr_in);

	if (listen(s, 5) != 0)
	{
		perror("listen");
		return -1;
	}
	if (debug)
		printf("listening.\n");

	if ( (accepted_sock = accept(s, (struct sockaddr *)&cli_addr, &cli_size)) == -1)
	{
		perror("accept");
		close(accepted_sock);
		return -1;
	}
	if (debug)
		printf("accepted to socket %d\n", accepted_sock);


	if (fork() != 0)
	{
		close(accepted_sock);
		return 1;
	}

	/* now i'm the child.  connect to remote_port and forward stuff back and
	 * forth */
	
	if ( (forward_sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)) == -1)
	{
		perror("socket");
		exit(-1);
	}
	if (debug)
		printf("got forwarding socket %d\n", forward_sock);
	
	if (connect(forward_sock, (struct sockaddr *) remote, sizeof(struct sockaddr_in)) == -1)
	{
		perror("connect");
		exit(-1);
	}
	
	if (flag)
	{
		snprintf(log_name, 80, "%s/%d_%s_%d", LOGS, local_port, remote_host, 
				remote_port);
		log = fopen(log_name, "w");
		fprintf(log, "\n---------------------------------------\n");
	}

	max_fds = (forward_sock > accepted_sock) ? forward_sock : accepted_sock;
	while (1)
	{
		FD_ZERO(&sockets);
		FD_SET(accepted_sock, &sockets);
		FD_SET(forward_sock, &sockets);
	
		bzero(buf, 1024);
		if (select(max_fds + 1, &sockets, NULL, NULL, NULL) == -1)
		{
			perror("select");
			exit(-1);
		}
		if (FD_ISSET(accepted_sock, &sockets))
		{
      if ( (n = recvfrom(accepted_sock, buf, 1024, 0, NULL, NULL)) < 1)
			{
				perror("recvfrom");
				exit(-1);
			}
			if (flag) fprintf(log, "%s\n-----\n", buf);
			if (send(forward_sock, buf, n, 0) == -1)
			{
				perror("send");
				exit(-1);
			}
		}
		else if (FD_ISSET(forward_sock, &sockets))
		{
			if ( (n = recvfrom(forward_sock, buf, 1024, 0, NULL, NULL)) < 1)
			{
				perror("recvfrom");
				exit(-1);
			}
			if (flag) fprintf(log, "%s\n-----\n", buf);
			if (send(accepted_sock, buf, n, 0) == -1)
			{
				perror("send");
				exit(-1);
			}
		}
		if (flag) fflush(log);
	}
	exit(0);
}

void sig_chld(int sig)
{
	pid_t pid;
	int s;
	pid = wait(&s);
}