/*************************************************************************
 ***      Author:  Rob Gubler -- tarsin@happybox.nu                    ***
 ***        Site:  www.hhp-programming.net                             ***
 ***        Date:  2001.06.22.r00                                      ***
 *** Description:  Basic TCP Proxy                                     ***
 ***    Comments:                                                      ***
 ***                                                                   ***
 *************************************************************************/


/*** Includes ***/
#include <stdio.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <errno.h>
#include <netinet/in.h>
#include <netdb.h>
#include <arpa/inet.h>
#include <sys/ioctl.h>

/*** Defines and Macros ***/

#define BACKLOG_LEN     5

/*** Structs, Typedefs, and Enums ***/

/*** Function Prototypes ***/
void ConnectClient(int cli_fd);
void RelayData(int srv_fd, int cli_fd);
int PrintHelp(char* cmd);

/*** Global Variables ***/
char* RemoteSrv;
int RemoteSrvPort;
int LocalPort;

char sbuf[2048];
char cbuf[2048];
extern int errno;

/*** Function Bodies ***/
int main(int argc, char* argv[]) {
    struct sockaddr_in cli_addr;
    struct sockaddr_in srv_addr;
    int cli_fd; 
    int cli_len; 
    int srv_fd;
    int opt = 1;
    int n, i;

    if(argc < 3 || argc > 4) {
        printf("\nError: Syntax Error.");
        return PrintHelp(argv[0]);
    } else {                                                /* set the ip and por
t to forward packets from/to */
        RemoteSrv = calloc(1, strlen(argv[1]) + 1);         
        memset(RemoteSrv, 0, (strlen(argv[1])+1));
        strcpy(RemoteSrv, argv[1]);
        
        RemoteSrvPort = atoi(argv[2]);
        if(argc > 3)
            LocalPort = atoi(argv[3]);
    }

    memset(&cli_addr, 0, sizeof(struct sockaddr_in));
    memset(&srv_addr, 0, sizeof(struct sockaddr_in));
    srv_addr.sin_family      = AF_INET;                     /* fill the sockaddr 
structure */
    srv_addr.sin_addr.s_addr = htonl(INADDR_ANY);
    srv_addr.sin_port        = htons(LocalPort);
    srv_fd = socket(PF_INET, SOCK_STREAM, 0);

    if( (n = bind(srv_fd, (struct sockaddr*)&srv_addr, sizeof(srv_addr))) == -1) 
{      /* bind to the local port */
        printf("Error: bind() returned %d", n);
            return 1;
    }
    listen(srv_fd, BACKLOG_LEN);                            /* listen for connect
ions */

    if(fork()) 
        return 1;

    for(i = 1; i;) {
        cli_len = sizeof(cli_addr);
            
        if((cli_fd = accept(srv_fd, (struct sockaddr *) &cli_addr, &cli_len)) < 0
) {    /* except the connection */
                if (errno == EINTR) 
                continue;
            else
                return 1;
            }

            switch(fork()) {
            case 0:                             /* child process  created*/
                close(srv_fd);                  /* close original socket */
                ConnectClient(cli_fd);                 /* process the request */
                close(cli_fd);
                return 0;
                break;
            case -1:                            /* child process couldn't be crea
ted */
                printf("\nError: fork() returned -1.  No child process could be c
reated.");
                break;
            default: 
                close(cli_fd);                  /* parent process */
                if (fork()) 
                return 1;
                break;
            }
    }
}

void ConnectClient(int cli_fd) {
    int dest;
    int found;
    struct sockaddr_in sa;
    struct hostent *hp;
    struct servent *sp;
    char gethost[128];
    char getport[128];
    char string[128];

        hp = gethostbyname(RemoteSrv);      /* get IP addr */
        if(hp) {                            /* if it resolved */
            found++;

        memcpy((caddr_t)&sa.sin_addr, hp->h_addr_list[0], hp->h_length);    /* ap
ply ip addr to sockaddr structure */
            memcpy((caddr_t)&sa.sin_addr, hp->h_addr, hp->h_length);

    } else {                            
            if (inet_addr(gethost) == -1) { /* if it did not resolve */
                found = 0;
                sprintf(string, "Error: Couldn't resolve the named address %s to 
an IP address\r\n", gethost);
                write(cli_fd, string, strlen(string));

        } else {                        /*  user iput was the decimal dotted nota
tion, not a named address */
                    found++;
                    sa.sin_addr.s_addr = inet_addr(gethost);
            }
        }
    sa.sin_family = AF_INET;            /* finish filling the structure */
    sa.sin_port = htons((unsigned) RemoteSrvPort);
    if (sa.sin_port == 0) {
        if((sp = getservbyname(getport, "tcp")))    /* verify the user assigned p
ort is valid */
                sa.sin_port = sp->s_port;
            else {
                sprintf(string, "Error:  %s: bad port number\r\n", getport);
                write(cli_fd, string, strlen(string));
                return;
            }
    }
    if ((dest = socket(AF_INET, SOCK_STREAM, 0)) < 0) {     /*  create the socket
 */
            perror("Error:  socket()  could not create socket");
            return;
    }
    connect(dest, (struct sockaddr *) &sa, sizeof(sa));     /* connect */
    fcntl(cli_fd, F_SETFL, O_NDELAY);
    fcntl(dest, F_SETFL, O_NDELAY);
    RelayData(dest, cli_fd);
    close(dest);
    return;
}

void RelayData(int srv_fd, int cli_fd) {
    char *chead, *ctail, *shead, *stail;
    int num, nfd, spos, cpos;
    extern int errno;
    fd_set rd, wr;

    chead = ctail = cbuf;
    cpos = 0;
    shead = stail = sbuf;
    spos = 0;
    while (1) {
        FD_ZERO(&rd);   /* initialize the sets */
        FD_ZERO(&wr);
        if(spos < sizeof(sbuf)-1) /* if we're ready to read more data */
        FD_SET(srv_fd, &rd);    /* check the read set */
        if(ctail > chead)
        FD_SET(srv_fd, &wr);    /* check write set */
        if(cpos < sizeof(cbuf)-1) 
        FD_SET(cli_fd, &rd);
        if(stail > shead) 
        FD_SET(cli_fd, &wr);
        nfd = select(256, &rd, &wr, 0, 0);  /*  get the next available ready set 
*/
        if(nfd <= 0) continue;          /* nothing ready. continue  */
        if(FD_ISSET(srv_fd, &rd)) {     /* read set returned for srv */
            num=read(srv_fd,stail,sizeof(sbuf)-spos);
            if ((num==-1) && (errno != EWOULDBLOCK)) return;
            if (num==0) return;
            if (num>0) {
                spos += num;
                stail += num;
                if (!--nfd) continue;
            }
        }
    if(FD_ISSET(cli_fd, &rd)) { /* read set returned for cli */
            num=read(cli_fd,ctail,sizeof(cbuf)-cpos);
            if ((num==-1) && (errno != EWOULDBLOCK)) return;
            if (num==0) return;
            if (num>0) {
                cpos += num;
                ctail += num;
                if (!--nfd) continue;
            }
        }
    if(FD_ISSET(srv_fd, &wr)) { /*  write set returned for srv */
            num=write(srv_fd,chead,ctail-chead);
            if ((num==-1) && (errno != EWOULDBLOCK)) return;
            if (num>0) {
                chead += num;
                if (chead == ctail) {
                    chead = ctail = cbuf;
                    cpos = 0;
                }
                if (!--nfd) continue;
            }
        }
    if(FD_ISSET(cli_fd, &wr)) { /* write set returned for cli */
            num=write(cli_fd,shead,stail-shead);
            if ((num==-1) && (errno != EWOULDBLOCK)) return;
            if (num>0) {
                    shead += num;
                    if (shead == stail) {
                        shead = stail = sbuf;
                        spos = 0;
                    }
                    if (!--nfd) 
                continue;
            }
        }
    }
}

int PrintHelp(char* cmd) {
    printf("\n\nSyntax: %s [Dest Addr] [Dest Port] [Local Port]", cmd);
    printf("\n\t[Dest Addr]  - IP or domain name that proxy will connect to.");
    printf("\n\t[Dest Port]  - The port the remote service is running on.");
    printf("\n\t[Local Port] - The port to bind locally that the proxy will run o
n. 1 - 65535");
    printf("\nExample: %s shell.server.com 23 80", cmd);
    printf("\n\n");

    return 1;

}