/*  ____       _   _       ____       ____       _   _ _        _    _ 
 * | __ )  ___| | | | __ _|  _ \ __ _| __ )  ___| | | | |__    / \  | |
 * |  _ \ / _ \ | | |/ _` | |_) / _` |  _ \ / _ \ | | | '_ \  / _ \ | |
 * | |_) |  __/ |_| | (_| |  _ < (_| | |_) |  __/ |_| | | | |/ ___ \|_|
 * |____/ \___|\___/ \__, |_| \_\__,_|____/ \___|\___/|_| |_/_/   \_(_)
 *                   |___/                                             
 */

/*
 * Tune me. This program will listen on LISTEN_PORT. If a client connects
 * and sends data, we redirect the traffic to DEFAULT_HOST:DEFAULT_PORT.
 * If a client connects but does not send any data for SSH_TIMEOUT ms,
 * we redirect the traffic to SSH_HOST:SSH_PORT. 
 */
#define LISTEN_PORT   80

#define DEFAULT_HOST  "localhost"
#define DEFAULT_PORT  8080

#define SSH_HOST      "localhost"
#define SSH_PORT      22

#define SSH_TIMEOUT   2000 /* milliseconds */

/*
 * Required headers
 */
#include <sys/types.h> /* socket, fork */
#include <sys/time.h> /* select */
#include <sys/socket.h> /* socket, inet_ntoa */
#include <netinet/in.h> /* htons, htonl, inet_ntoa */
#include <arpa/inet.h> /* inet_ntoa */
#include <netdb.h> /* gethostbyname */
#include <unistd.h> /* fork */
#include <signal.h> /* signal */
#include <stdlib.h> /* realloc, free */
#include <stdio.h> /* stderr */
#include <string.h> /* strcmp */
#include <errno.h> /* errno */

/*
 * You should not need to change this.
 */
#define BUFFER_SIZE 1024
#define MAX_CLIENTS 20

/*
 * Prototypes
 */
static int  Connect ( char *, int );
static void Listen  ( int );
static int  Copy    ( int, int );

/*
 * The main loop
 */
int main( int i_argc, char *ppsz_argv[] )
{
    int i_ret;
    int i_dummy;
    fd_set fdset;
    struct timeval timer;

    /* Server information */
    int i_fd;
    struct sockaddr_in server;

    /* Client information */
    int i_newfd;
    struct sockaddr_in client;
    socklen_t addrlen;

    /* Ignore SIG_CHLD to avoid zombies */
    signal( SIGCHLD, SIG_IGN );

    /* Create socket */
    i_fd = socket( PF_INET, SOCK_STREAM, 0 );
    if( i_fd < 0 )
    {
        fprintf( stderr, "Failed to create socket\n" );
        return -1;
    }

    /* Configure socket to reuse address */
    i_dummy = 1;
    i_ret = setsockopt( i_fd, SOL_SOCKET, SO_REUSEADDR, (void *) &i_dummy,
                        sizeof( i_dummy ) );
    if( i_ret == -1 )
    {
        fprintf( stderr, "Failed to setsockopt SO_REUSEADDR (%s)\n",
                         strerror(errno) );
        return -1;
    }

    /* Set remaining socket information */
    memset( &server, 0x00, sizeof( struct sockaddr_in ) );
    server.sin_family = AF_INET;
    server.sin_port = htons( LISTEN_PORT );
    server.sin_addr.s_addr = htonl( INADDR_ANY );

    /* Bind socket */
    i_ret = bind( i_fd, (struct sockaddr *) &server, sizeof( server ) );
    if( i_ret < 0 )
    {
        fprintf( stderr, "Failed to bind socket (%s)\n", strerror(errno) );
        return -1;
    }

    /* Set socket to listen mode */
    i_ret = listen( i_fd, MAX_CLIENTS );
    if( i_ret < 0 )
    {
        fprintf( stderr, "Failed to listen to socket (%s)\n", strerror(errno) );
        return -1;
    }

    fprintf( stderr, "Main loop listening\n" );

    /* Loop until we get a connection */
    do
    {
        FD_ZERO( &fdset );
        FD_SET( i_fd, &fdset );

        timer.tv_sec = 0;
        timer.tv_usec = 10000;

        /* Wait for a connection */
        if( select( i_fd + 1, &fdset, NULL, NULL, &timer ) )
        {
            addrlen = sizeof( struct sockaddr );
            i_newfd = accept( i_fd, (struct sockaddr *) &client, &addrlen );
            if( i_newfd == -1 )
            {
                fprintf( stderr, "Failed accepting connection\n" );
                continue;
            }

            fprintf( stderr, "New client %s\n", inet_ntoa( client.sin_addr ) );

            switch( fork() )
            {
                /* We are the child, handle the socket */
                case 0:
                    Listen( i_newfd );
                    return 0;
                    break;
    
                /* Boo. */
                case -1:
                    fprintf( stderr, "Fork failed\n" );
                    break;

                /* We are the parent. Just wait for data... */
                default:
                    if( close( i_newfd ) )
                    {
                        fprintf( stderr, "Failed to close connection %d (%s)\n",
                                         i_newfd, strerror(errno) );
                    }
                    break;
            }
        }

    } while( 1 );

    return 0;
}

/*
 * Listen to what the client says
 */
void Listen( int i_client )
{
    /* Descriptor set and timer for select */
    fd_set fdset;
    int i_select = i_client + 1;
    struct timeval timer;
    int b_alive = 1;

    /* The server */
    int i_server = -1;

    /* Timeout for protocol guessing */
    int i_timeout = SSH_TIMEOUT / 10;
    enum { NONE, SSH, DEFAULT } i_proto = NONE;

    do
    {
        /* See whether it's a client or a server */
        if( i_proto == NONE )
        {
            i_timeout--;
            if( i_timeout == 0 )
            {
                fprintf( stderr, "No data for %i milliseconds, assuming SSH\n",
                                 SSH_TIMEOUT );
                i_proto = SSH;
            }
        }

        /* Initialize descriptor set */
        FD_ZERO( &fdset );
        FD_SET( i_client, &fdset );

        if( i_server == -1 )
        {
            switch( i_proto )
            {
                case NONE:
                    break;
                case SSH:
                    i_server = Connect( SSH_HOST, SSH_PORT );
                    if( i_server > i_client ) i_select = i_server + 1;
                    break;
                case DEFAULT:
                    i_server = Connect( DEFAULT_HOST, DEFAULT_PORT );
                    if( i_server > i_client ) i_select = i_server + 1;
                    break;
            }
        }
        else
        {
            FD_SET( i_server, &fdset );
        }

        /* Initialize timer */
        timer.tv_sec = 0;
        timer.tv_usec = 10000;

        /* Wait for data */
        if( select( i_select, &fdset, NULL, NULL, &timer ) )
        {
            if( i_proto == NONE )
            {
                if( ! FD_ISSET( i_client, &fdset ) )
                {
                    continue;
                }

                /* We got data, which means it's not an SSH client */
                fprintf( stderr, "Got data, using default protocol\n" );
                i_proto = DEFAULT;
            }

            /* Eeek, we aren't connected yet! */
            if( i_server == -1 )
            {
                continue;
            }

            if( FD_ISSET( i_client, &fdset ) )
            {
                if( Copy( i_server, i_client ) <= 0 )
                {
                    close( i_client );
                    close( i_server );
                    return;
                }
            }

            if( FD_ISSET( i_server, &fdset ) )
            {
                if( Copy( i_client, i_server ) <= 0 )
                {
                    close( i_client );
                    close( i_server );
                    return;
                }
            }
        }
    }
    while( b_alive );

    fprintf( stderr, "Closing connection %i\n", i_client );
    close( i_client );
}

/*
 * Connect to a remote server
 */
static int Connect( char *psz_server, int i_port )
{
    int i_ret;

    /* Socket information */
    int i_fd;
    struct sockaddr_in server;

    /* Create socket */
    i_fd = socket( AF_INET, SOCK_STREAM, 0 );
    if( i_fd < 0 )
    {
        fprintf( stderr, "Failed to create socket (%s)\n", strerror(errno) );
        return -1;
    }

    /* Set remaining socket information */
    memset( &server, 0x00, sizeof( struct sockaddr_in ) );
    server.sin_family = AF_INET;
    server.sin_port = htons( i_port );
    if(!inet_aton(psz_server, &server.sin_addr))
    {
        struct hostent* p_hostent = gethostbyname(psz_server);
        if(p_hostent == NULL)
        {
            fprintf( stderr, "Invalid server name %s\n", psz_server );
            return -1;
        }

        memcpy(&server.sin_addr, p_hostent->h_addr, p_hostent->h_length);
    }

    /* Bind socket */
    i_ret = connect( i_fd, (struct sockaddr *) &server, sizeof( server ) );
    if( i_ret < 0 )
    {
        fprintf( stderr, "Failed to connect socket\n" );
        return -1;
    }

    return i_fd;
}

/*
 * Copy data from one fd to another.
 */
int Copy( int i_to, int i_from )
{
    /* Buffer to store read data */
    static unsigned char p_buffer[BUFFER_SIZE];

    int i_read = read( i_from, p_buffer, BUFFER_SIZE );

    if( i_read < 0 )
    {
        fprintf( stderr, "Error while reading from client (%s)\n",
                         strerror(errno) );
        return i_read;
    }
    else if( i_read == 0 )
    {
        fprintf( stderr, "Connection closed by peer\n" );
        return i_read;
    }

    write( i_to, p_buffer, i_read );

    return i_read;
}