/*
 * Fully featured multi threaded Tcp relayer. This code is copyright.
 * Usage is subject to giving credit to the original author. Otherwise
 * it may not be used. It may not be published on any site other than 
 * the one it was found without prior permission from author.
 *
 * Author: Hi_Tech_Assassin
 *
 */

#include <windows.h>
#include <stdio.h>

#pragma comment (lib,"ws2_32") 

#define TIMEOUT 2
#define BUFFLEN	10000

class redirect
{
	friend class connection;
	
	// our connection class needs access to 
	// certain definitions belonging to this class

	public:
		
		void start();
		void setLocalPort(int port);
		void setRemotePort(int port);
		void setRemoteHost(const char *host);

	private:

		int DoWinsock();
		void SpawnThread();
		SOCKET SetUpListener();
		void AcceptConnections(SOCKET ListeningSocket);

		int *REMOTEPORT; 
		int *LOCALPORT;
		char *REMOTEHOST;
};

redirect redirector;

// this class is needed for each connection

class connection
{
	friend class redirect;

	private:

		SOCKET ServerFd;

		void ServerThreadStart()
		{
			DWORD ThreadID;
			CreateThread(NULL, 0,ServerDispatcher, this, 0, &ThreadID);
		}

		static DWORD WINAPI ClientDispatcher(void* p)
		{ 
			return reinterpret_cast<connection*>(p)->ClientHandle(); 
		}

		static DWORD WINAPI ServerDispatcher(void* p)
		{
			return reinterpret_cast<connection*>(p)->ServerHandle(); 
		}

		void ClientThreadStart()
		{	
                        DWORD ThreadID;
                        CreateThread(NULL, 0,ClientDispatcher,this, 0, &ThreadID);
		}

		DWORD ServerHandle();
		DWORD ClientHandle();
		int InitClient();
		SOCKET ClientFd; 
};

void redirect::setLocalPort(int port)
{
	this->LOCALPORT=(int*)port;
}

void redirect::setRemotePort(int port)
{
	this->REMOTEPORT=(int*)port;
}

void redirect::setRemoteHost(const char *host)
{
	this->REMOTEHOST=(char*)host;
}

void redirect::start()
{
	this->DoWinsock();
}

int redirect::DoWinsock()
{
    printf("Establishing the listener...\n");
    SOCKET ListeningSocket;

    ListeningSocket=this->SetUpListener();

    if (ListeningSocket == INVALID_SOCKET) 
    {
        printf("Invalid socket\n");
        return 3;
    }

    printf("Waiting for connections...\n");

    while (1) 
    {
        this->AcceptConnections(ListeningSocket);
        printf("Acceptor restarting\n");
    }
}

SOCKET redirect::SetUpListener()
{
    SOCKET sd = socket(AF_INET, SOCK_STREAM, 0);

    if (sd != INVALID_SOCKET) 
    {
            sockaddr_in sinInterface;
            sinInterface.sin_family = AF_INET;
            sinInterface.sin_addr.s_addr = htonl(INADDR_ANY);
            sinInterface.sin_port = htons((unsigned short)this->LOCALPORT);

            if (bind(sd, (sockaddr*)&sinInterface,sizeof(sockaddr_in)) != SOCKET_ERROR) 
            {
                listen(sd, SOMAXCONN);
                return sd;
            }
            else printf("bind() failed");
    }

    return INVALID_SOCKET;
}


void redirect::AcceptConnections(SOCKET ListeningSocket)
{
    sockaddr_in sinRemote;
    int nAddrSize = sizeof(sinRemote);

    while (1)
    {
        SOCKET sd = accept(ListeningSocket, 
			(sockaddr*)&sinRemote, &nAddrSize);

        if (sd != INVALID_SOCKET) 
        {
    	    printf("Accepted connection from %s:%d..\n",inet_ntoa(sinRemote.sin_addr),
            ntohs(sinRemote.sin_port));
            connection con; 
            con.ServerFd=sd;
            con.ServerThreadStart();
        }
        else 
        {
            printf("accept() failed..\n");
            return;
        }
    }
}

DWORD connection::ServerHandle()
{
	int nReadBytes;
	char received[BUFFLEN];

	if(!InitClient())
	{
		puts("Remote host unreachable or timed out\n");
		closesocket((unsigned int)this->ServerFd); 
		exit(-1);
	}

	this->ClientThreadStart();

	while((nReadBytes = recv((unsigned int)this->ServerFd, received, 
		sizeof(received)-1, 0))!= SOCKET_ERROR && nReadBytes!=0)
	{
		received[nReadBytes]='\0';

		if(send((unsigned int)this->ClientFd, received, nReadBytes, 0)==-1)
			break;

	}

	closesocket((unsigned int)this->ClientFd); 

	printf("Party bailed out\n");

	return 0;
}

int connection::InitClient()
{
	struct hostent 		*pTarget;
	struct sockaddr_in 	sock;

	SOCKET sd = socket(AF_INET, SOCK_STREAM, 0);

	this->ClientFd=sd;

	if(sd==INVALID_SOCKET)
		return 0;

	if ((pTarget = gethostbyname(redirector.REMOTEHOST)) == NULL)
		return 0;

	memcpy(&sock.sin_addr.s_addr, pTarget->h_addr, pTarget->h_length);
	sock.sin_family = AF_INET;
	sock.sin_port = htons((USHORT)redirector.REMOTEPORT);

	struct timeval timeout;
	unsigned long y = 1;

	ioctlsocket(sd,FIONBIO,&y);

	timeout.tv_sec=TIMEOUT;
	timeout.tv_usec = 0;

	if((connect(sd, (struct sockaddr *)&sock, sizeof (sock) ))==-1)
	{
		fd_set writefds;
		fd_set exceptfds;

		FD_ZERO (&writefds);
		FD_ZERO (&exceptfds);
		FD_SET (sd, &writefds);
		FD_SET (sd, &exceptfds);

		select (0, NULL, &writefds, &exceptfds, &timeout); 

		if (!FD_ISSET (sd, &writefds))
		{
			closesocket(sd);
			return 0;
		}
		else
		{
			y=0;
			ioctlsocket(sd,FIONBIO,&y);
		}
	}

	return sd;
}

DWORD connection::ClientHandle()
{
	int nReadBytes;
	char received[BUFFLEN];

	while((nReadBytes = recv((unsigned int)this->ClientFd, received, 
	sizeof(received)-1, 0))!= SOCKET_ERROR && nReadBytes!=0)
	{
		received[nReadBytes]='\0';

		if(send((unsigned int)this->ServerFd, received, nReadBytes, 0)==-1)
			break;
	}

	closesocket((unsigned int)this->ServerFd);
	return 0;
}

int main(int argc, char* argv[])
{ 
        WSAData wsaData;
        int nCode;

        if(argc!=4)
        {
               printf("\nUsage:\n %s <remote host> <localport> <remoteport>\n\n", argv[0]);
               return 0;
        }

    	if ((nCode = WSAStartup(MAKEWORD(1, 1), &wsaData)) != 0) 
    	{
               printf("WSAStartup() returned error code %d", nCode);
               return 255;
    	}

        const char* remotehost=argv[1];
        const int localport=atoi(argv[2]);
        const int remoteport=atoi(argv[3]);

        redirector.setLocalPort(localport);
        redirector.setRemoteHost(remotehost);
        redirector.setRemotePort(remoteport);
        redirector.start();

    	WSACleanup();
    	return 0;
}