/*  Author: David Foster      Date Created: June 18, 2000
 *  Last Modified: February 26, 2001
 *  DESCRIPTION: Port forwarder.  Provide the program a local port and 
 *	the host and port you want the local port forwarded to.  All 
 *	connections to the chosen port on the localhost will be forwarded
 *	to the chosen server:port.  Both TCP and UDP connections can 
 *	be used, but not at the same time. 
 *  [rservadd]<----->[lservaddr]==(forward daemon)==[lcliaddr]<----->[rcliaddr]
 *  See the print_usage function below for more explaination
 *  TO DO: need to fix udp part, right now using the udp parameter
 *	causes uncontrollable forking that loads down the CPU
 *	needless to say, don't use the udp option
*/

#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h> 	 // ditto  MIGHT NOT NEED THIS, WE DON"T USE struct timeval.....
#include <netinet/in.h>
#include <arpa/inet.h>  // inet_pton, inet_ntop, inet_aton, inet_addr, inet_ntoa
#include <errno.h>
#include <netdb.h>
#include <stdio.h>
#include <stdlib.h>  //for atoi
#include <string.h>  //for strcmp
#include <unistd.h>	 // ditto (for fdset struct anyway)
#include <sys/select.h>  // for select system call
#include <signal.h>	//for killing zombie children

#define maxfd 3
#define MAXSOCKADDR 128
#ifndef INET_ADDRSTRLEN
#define INET_ADDRSTRLEN 16
#endif

// Function Prototypes
void sys_err(char* str);
void print_usage();
void zombie_killer(int signo);


int main(int argc, char **argv){
    int i, type_array[2];
    pid_t pid;
    char line[4096], str_addr[INET_ADDRSTRLEN], *proto_array[2];
    char raddr_info[INET_ADDRSTRLEN], laddr_info[INET_ADDRSTRLEN];
    int rport_info, lport_info;
    struct sockaddr_in  rservaddr,// remote server connection address info 
			lservaddr,// local server connection address info
			rcliaddr, // remote client connection address info
			lcliaddr; // local client connection address info
    socklen_t len;
    struct servent *sserv;
    struct hostent *shost;
    int numfd, 
	servdoneflag=0, 
	clidoneflag=0, 
	listensock, 
	ssock, 
	connectsock,
	protocol = 0, 
	localport, 
	remoteport;
    fd_set readset;
	bzero(&rservaddr, sizeof(rservaddr));
	bzero(&lservaddr, sizeof(lservaddr));
	bzero(&rcliaddr, sizeof(rcliaddr));
	bzero(&lcliaddr, sizeof(lcliaddr));
	rservaddr.sin_family = lservaddr.sin_family = AF_INET;
	rcliaddr.sin_family = lcliaddr.sin_family = AF_INET;
	lcliaddr.sin_addr.s_addr = htonl(INADDR_ANY);
	type_array[0] = SOCK_STREAM; type_array[1] = SOCK_DGRAM;
	proto_array[0] = "tcp"; 
	proto_array[1] = "udp";
	signal(SIGCHLD, zombie_killer);
	if(argc < 4) print_usage();

/////////////////////////////////////////////////////////////
// Convert the host argument to the proper format
/////////////////////////////////////////////////////////////
	if(inet_pton(AF_INET, argv[2], &rservaddr.sin_addr) <= 0){
//	    fprintf(stderr, "inet_pton didn't work\n");
	    printf("Attempting to resolve hostname %s---->", argv[2]);
	    fflush(NULL);
	    if((shost = gethostbyname(argv[2])) == NULL){
	       perror("Could not resolve hostname");
	       exit(0);
	    }else{
	       memcpy(&rservaddr.sin_addr, (struct in_addr *) shost->h_addr, sizeof(struct in_addr));
	       inet_ntop(AF_INET, &rservaddr.sin_addr, str_addr, INET_ADDRSTRLEN); 
	       printf("%s\n", str_addr);
	    }
	}

//////////////////////////////////////////////////////////////
// Check all the arguments for validity
// (Some of this checking could be better....)
//////////////////////////////////////////////////////////////
	localport = atoi(argv[1]);
	if(localport < 1024 || localport > 65535) 
		sys_err("Illegal local port number.  Must be 1024<PORT<65535.\n\n");
	remoteport = atoi(argv[3]);
	if(remoteport < 1 || remoteport > 65535) 
		sys_err("Illegal remote port number.  Must be 1<PORT<65535.\n\n");
	if(argc == 5){  // We've got a protocol argument to deal with
	    if((strcmp(argv[4], "tcp")) == 0){      
		protocol = 0;
	    } else if((strcmp(argv[4], "udp")) == 0){ 
		protocol = 1;
		sys_err("udp protocol currently not available in this program.\n\n");
	    } else{ 
		sys_err("Illegal protocol.  Must be either tcp or udp.\n");
	    }
	}

	printf("Local Port: %d  Remote Address: %s  Remote Port: %d  Protocol: %s\n", 
		localport, str_addr, remoteport, proto_array[protocol]);

/////////////////////////////////////////////////////////////////////	
// Check to see if remote server is available
////////////////////////////////////////////////////////////////////
	rservaddr.sin_port = htons(remoteport);
	if((ssock = socket(AF_INET, type_array[protocol], 0)) < 0){
		perror("error creating socket");
		exit(0);
	}
	printf("Checking for remote server availability....."); 
	fflush(NULL);
	if(connect(ssock, (struct sockaddr*) &rservaddr, sizeof(rservaddr)) == 0){
	   printf("Server active!\n");
	   shutdown(ssock, 2);
	}else{
	    sys_err("Connect error!\nRemote server apparently not available!\n\n"); 
	}
	close(ssock);

/////////////////////////////////////////////////////////////////
// Check to see if we can bind the listening server to the 
// local port (any interface)
/////////////////////////////////////////////////////////////////
	printf("Attempting to bind forwarding daemon to local port %d....", localport);
	fflush(NULL);
	lservaddr.sin_port = htons(localport);
	if((listensock = socket(AF_INET, type_array[protocol], 0)) < 0){
		perror("error creating socket");exit(0);
	}
	if(bind(listensock, (struct sockaddr*) &lservaddr, sizeof(lservaddr)) >= 0){
	    printf("Ok!\nListening for incoming connections on all interfaces\n"); 
	    fflush(NULL);
	    listen(listensock, 5);
	}else{
	    sys_err("Bind error!\n"); 
	}
	i = sizeof(rcliaddr);

/////////////////////////////////////////////////////////////////
// We need to wait for a client.  When one arrives, fork and process
// the client (connect to server and start forwarding packets about),
// while server continues to listen for new clients.
/////////////////////////////////////////////////////////////////
	while(1){
	  connectsock = accept(listensock, (struct sockaddr*)&rcliaddr, &i);
	  if((pid = fork()) == 0){
	    close(listensock);
	    if((ssock = socket(AF_INET, type_array[protocol], 0)) < 0){
		perror("error creating socket");exit(0);
	    }
	    if(connect(ssock, (struct sockaddr*) &rservaddr, sizeof(rservaddr)) < 0) sys_err("Connect error!\nRemote server apparently not available!\n\n"); 

///////////////////////////////////////////////////////////////////
//We got two connections going now, one to client and one to server.  Yay!
///////////////////////////////////////////////////////////////////
	    len = sizeof(lcliaddr);
	    getpeername(connectsock, (struct sockaddr*)&lcliaddr, &len); 
	    inet_ntop(AF_INET, &lcliaddr.sin_addr, raddr_info,INET_ADDRSTRLEN);
	    rport_info = ntohs(lcliaddr.sin_port);
	    printf("New Connection:\n\t %s:%d<-->[%d:localhost:", 
		raddr_info, rport_info, remoteport);
	    len = sizeof(lservaddr);
	    getsockname(ssock, (struct sockaddr*)&lcliaddr, &len);
	    inet_ntop(AF_INET, &lcliaddr.sin_addr, laddr_info,INET_ADDRSTRLEN);
	    lport_info = ntohs(lcliaddr.sin_port);
	    printf(":%d]<-->%s:%d\n", 
		localport, laddr_info, lport_info);

	    FD_ZERO(&readset);
	    FD_SET(ssock, &readset);
	    FD_SET(connectsock, &readset);
//	    printf("connectsock: %d  ssock: %d\n", connectsock, ssock);
	    while((!servdoneflag) && (!clidoneflag)){
	      FD_SET(ssock, &readset);
	      FD_SET(connectsock, &readset);
	      numfd = select(FD_SETSIZE, &readset, NULL, NULL, NULL);
	      if(FD_ISSET(connectsock, &readset)){
		if( (i = read(connectsock, line, 4096)) == 0){  // connection closed
		    clidoneflag = 1;
		}else{
		    if(write(ssock, line, i) != i){
			printf("Insufficient data written to server\n");
		    }
		}
	      } 
	      if(FD_ISSET(ssock, &readset)){
		  if( (i = read(ssock, line, 4096)) == 0){ //connection closed
		    servdoneflag = 1;
		    shutdown(servdoneflag, 2);
		  }else{
		    if(write(connectsock, line, i) != i){
			printf("Insufficient data written to client\n");
		    }
		  }
	      }
	    }//end child forwarding while loop
	    shutdown(ssock, 2);
	    shutdown(connectsock, 2);
	    exit(0);
	  }//end child process fork
	  close(connectsock);
	}//end infinite while loop

	close(listensock);
	exit(0);
}
 

void sys_err(char* str){
	fprintf(stderr, "%s",  str);
	exit(1);
	return;
}

void print_usage(){
	fprintf(stderr, "\nUsage: forward <local port> <remote host> <remote port> [protocol]\n");
	fprintf(stderr, "\t[protocol] value must be either tcp (default) or udp\n\n");
	fprintf(stderr, "Forward is a port forwarding program that takes incoming connections to\n");
	fprintf(stderr, "the localhost:localport and forwards it to the remotehost:remoteport\n");
	fprintf(stderr, "and vice-versa, simply acting as a middleman for a client-server\n\n");
	exit(1);
	return;
}

void zombie_killer(int signo){
	int stat;
	wait(&stat);
	return;
}