#include <iostream>
#include <boost/asio.hpp>
#include <boost/lexical_cast.hpp>
#include <thread>
#include <array>
#include <iomanip>

#include <boost/bind.hpp>

#define SSL_R_SHORT_READ 219
#include "ssl/ssl_locl.h"
#include <boost/asio/ssl.hpp>


#if defined(close)
#undef close
#endif

enum AddressType
{
	AT_IPV4 = 0,
	AT_HOST = 3,
	AT_IPV6 = 4
};

class ConnectResponseRecord
{
public:
	std::vector<unsigned char> rawData;

	enum ConnectResponseType
	{
		CRT_SUCCESS = 0,
		CRT_GENERAL_SOCKS_SERVER_FAILURE = 1,
		CRT_CONNECTION_NOT_ALLOWED_BY_RULESET = 2,
		CRT_NETWORK_UNREACHABLE = 3,
		CRT_HOST_UNREACHABLE = 4,
		CRT_CONNECTION_REFUSED = 5,
		CRT_TTL_EXPIRED = 6,
		CRT_COMMAND_NOT_SUPPORTED = 7,
		CRT_ADDRESS_TYPE_NOT_SUPPORTED = 8,
		CRT_TO_FF_UNASSIGNED = 9
	};

	ConnectResponseType connectResponseType;

	AddressType addressType;

	std::string address;

	uint16_t port;
};

class ConnectRequestRecord
{
public:
	std::string address;
	uint16_t port;

	std::string getRequestData();
};

std::string ConnectRequestRecord::getRequestData()
{
	std::string data = "";
	data += 0x05;
	data += 0x01;
	data += char(0x00);
	data += 0x03;
	data += static_cast<unsigned char>(address.size());
	data += address;
	data += port / 256;
	data += port % 256;
	return data;
}




typedef boost::asio::ssl::stream<boost::asio::ip::tcp::socket> ssl_socket;


class ProxySession :
	public std::enable_shared_from_this<ProxySession>
{
public:
	ProxySession(std::shared_ptr<ssl_socket> socket, boost::asio::io_service& inIoService)
		: ioService(inIoService)
		, mSocket(socket)
		, outsideConnectSocket(ioService)
	{
		std::cout << "ProxySession Create" << std::endl;
	}

	~ProxySession()
	{
		std::cout << "ProxySession Destroy" << std::endl;
	}

	void start()
	{
		std::cout << "ProxySession::start" << std::endl;
		asyncHandshake();
		//readClientVersion();
	}


private:

	ssl_socket& socket()
	{
		return *mSocket;
	}

	ssl_socket::lowest_layer_type& lowerSocket()
	{
		return mSocket->lowest_layer();
	}


	void asyncHandshake()
	{
		std::cout << "ProxySession::asyncHandshake" << std::endl;
		auto self(shared_from_this());

		socket().async_handshake(boost::asio::ssl::stream_base::server,
			[this, self](boost::system::error_code ec) {
		
			std::cout << "ProxySession::asyncHandshake inner" << std::endl;

			if (!ec)
			{
				readClientVersion();
			}
			else
			{
				lowerSocket().close();
			}
		
		});
	}

	/*
	void handleHandshake(const boost::system::error_code& error)
	{
		if (!error)
		{
			readClientVersion();
		}
		else
		{
			lowerSocket().close();
		}
	}*/


	std::array<unsigned char, 3> clientVersion;
	
	void readClientVersion()
	{
		std::cout << "ProxySession::readClientVersion" << std::endl;
		auto self(shared_from_this());

		boost::asio::async_read(socket(),
			boost::asio::buffer(clientVersion.data(), clientVersion.size()),
			[this, self](boost::system::error_code ec, std::size_t /*length*/)
		{
			std::cout << "ProxySession::readClientVersion inner" << std::endl;
			if (!ec)
			{
				if (clientVersion[0] == 0x05 && clientVersion[1] == 0x01 && clientVersion[2] == 0x02)
				{
					sendServerVersion();
				}
				else
				{
					lowerSocket().close();
				}
			}
			else
			{
				lowerSocket().close();
			}
		});
	}


	void sendServerVersion()
	{

		std::cout << "ProxySession::sendServerVersion" << std::endl;

		auto self(shared_from_this());
	
		std::array<char, 2> version = { 0x05, 0x02 };

		boost::asio::async_write(socket(),
			boost::asio::buffer(version.data(), version.size()),
			[this, self](boost::system::error_code ec, std::size_t length)
		{
			std::cout << "ProxySession::sendServerVersion inner" << std::endl;
			if (!ec)
			{
				readLoginPassword();
			}
			else
			{
				lowerSocket().close();
			}
		});
	}


	void readLoginPassword()
	{
		std::string login = "telegram-proxy-user";
		std::string password = "telegram-telegram-999112";
		std::string expectedData = "";
		expectedData += 0x01;
		expectedData += static_cast<unsigned char>(login.size());
		expectedData += login;
		expectedData += static_cast<unsigned char>(password.size());
		expectedData += password;

		std::shared_ptr<std::string> clientLoginPasswordPtr = std::make_shared<std::string>();


		clientLoginPasswordPtr->resize(expectedData.size());


		auto self(shared_from_this());


		boost::asio::async_read(socket(),
			boost::asio::buffer(&((*clientLoginPasswordPtr)[0]), clientLoginPasswordPtr->size()),
			[this, self, expectedData, clientLoginPasswordPtr](boost::system::error_code ec, std::size_t /*length*/)
		{
			if (!ec)
			{
				if (*clientLoginPasswordPtr == expectedData)
				{
					sendAuthStatus();
				}
				else
				{
					lowerSocket().close();
				}
			}
			else
			{
				lowerSocket().close();
			}
		});

	}

	std::array<unsigned char, 2> authStatus = { 0x01, 0x00 };

	void sendAuthStatus()
	{
		auto self(shared_from_this());

		

		boost::asio::async_write(socket(),
			boost::asio::buffer(authStatus.data(), authStatus.size()),
			[this, self](boost::system::error_code ec, std::size_t length)
		{

			if (!ec)
			{
				readConnectRequest();
			}
			else
			{
				lowerSocket().close();
			}
		});
	}

	void readConnectRequest()
	{
		std::shared_ptr<std::array<unsigned char, 5>> firstPartPtr = std::make_shared<std::array<unsigned char, 5>>();


		auto self(shared_from_this());


		boost::asio::async_read(socket(),
			boost::asio::buffer(firstPartPtr->data(), firstPartPtr->size()),
			[this, self, firstPartPtr](boost::system::error_code ec, std::size_t /*length*/)
		{
			if (!ec)
			{
				if ((*firstPartPtr)[0] == 0x05 && (*firstPartPtr)[1] == 0x01 && (*firstPartPtr)[2] == 0x00)
				{
					if ((*firstPartPtr)[3] == 0x03)
					{
						readConnectRequestPart2(firstPartPtr);
					}
					else if ((*firstPartPtr)[3] == 0x01)
					{
						readConnectRequestPart2IpAddress(firstPartPtr);
					}

				}
				else
				{
					lowerSocket().close();
				}
			}
			else
			{
				lowerSocket().close();
			}
		});

	}


	void readConnectRequestPart2(std::shared_ptr<std::array<unsigned char, 5>> firstPartPtr)
	{
		unsigned int len = (*firstPartPtr)[4];

		std::shared_ptr<std::vector<unsigned char>> secondPartPtr = std::make_shared<std::vector<unsigned char>>();


		secondPartPtr->resize(len + 2);

		auto self(shared_from_this());


		boost::asio::async_read(socket(),
			boost::asio::buffer(secondPartPtr->data(), secondPartPtr->size()),
			[this, self, firstPartPtr, secondPartPtr, len](boost::system::error_code ec, std::size_t /*length*/)
		{
			if (!ec)
			{

				ConnectRequestRecord connectRequestRecord;

				connectRequestRecord.address = std::string(&((*secondPartPtr)[0]), &((*secondPartPtr)[0]) + len);
				connectRequestRecord.port = (*secondPartPtr)[len] * 256 + (*secondPartPtr)[len + 1];

				boost::asio::ip::tcp::resolver resolver(this->ioService);
				auto endpointIterator = resolver.resolve({ connectRequestRecord.address, boost::lexical_cast<std::string>(connectRequestRecord.port) });


				doConnectOutput(endpointIterator, connectRequestRecord);

			}
			else
			{
				lowerSocket().close();
			}
		});

	}


	void readConnectRequestPart2IpAddress(std::shared_ptr<std::array<unsigned char, 5>> firstPartPtr)
	{
		unsigned int ip1st = (*firstPartPtr)[4];

		std::shared_ptr<std::array<unsigned char, 5>> secondPartPtr = std::make_shared<std::array<unsigned char, 5>>();


		auto self(shared_from_this());


		boost::asio::async_read(socket(),
			boost::asio::buffer(secondPartPtr->data(), secondPartPtr->size()),
			[this, self, firstPartPtr, secondPartPtr, ip1st](boost::system::error_code ec, std::size_t /*length*/)
		{
			if (!ec)
			{

				ConnectRequestRecord connectRequestRecord;

				connectRequestRecord.address = boost::lexical_cast<std::string>(static_cast<unsigned int>(ip1st)) + "."
					+ boost::lexical_cast<std::string>(static_cast<unsigned int>((*secondPartPtr)[0])) + "."
					+ boost::lexical_cast<std::string>(static_cast<unsigned int>((*secondPartPtr)[1])) + "."
					+ boost::lexical_cast<std::string>(static_cast<unsigned int>((*secondPartPtr)[2]));

				connectRequestRecord.port = (*secondPartPtr)[3] * 256 + (*secondPartPtr)[4];

				boost::asio::ip::tcp::resolver resolver(this->ioService);
				auto endpointIterator = resolver.resolve({ connectRequestRecord.address, boost::lexical_cast<std::string>(connectRequestRecord.port) });


				doConnectOutputIpAddress(endpointIterator, connectRequestRecord, { static_cast<unsigned char>(ip1st) , ((*secondPartPtr)[0]) , ((*secondPartPtr)[1]) , ((*secondPartPtr)[2]) });

			}
			else
			{
				lowerSocket().close();
			}
		});

	}


	void doConnectOutput(boost::asio::ip::tcp::resolver::iterator endpointIterator, ConnectRequestRecord connectRequestRecord)
	{
		auto self(shared_from_this());

		boost::asio::async_connect(outsideConnectSocket, endpointIterator,
			[this, self, connectRequestRecord](boost::system::error_code ec, boost::asio::ip::tcp::resolver::iterator)
		{
			if (!ec)
			{
				sendConnectResponse(connectRequestRecord);
			}
			else
			{
				lowerSocket().close();
				outsideConnectSocket.close();
			}
		});
	}

	void doConnectOutputIpAddress(boost::asio::ip::tcp::resolver::iterator endpointIterator, ConnectRequestRecord connectRequestRecord, std::array<unsigned char, 4> ipAddress)
	{
		auto self(shared_from_this());

		boost::asio::async_connect(outsideConnectSocket, endpointIterator,
			[this, self, connectRequestRecord, ipAddress](boost::system::error_code ec, boost::asio::ip::tcp::resolver::iterator)
		{
			if (!ec)
			{
				sendConnectResponseIpAddress(connectRequestRecord, ipAddress);
			}
			else
			{
				lowerSocket().close();
				outsideConnectSocket.close();
			}
		});
	}


	void sendConnectResponse(ConnectRequestRecord connectRequestRecord)
	{
		auto self(shared_from_this());

		std::string connectResponse;

		connectResponse += 0x05;
		connectResponse += static_cast<char>(0x00);
		connectResponse += static_cast<char>(0x00);
		connectResponse += 0x03;

		connectResponse += static_cast<unsigned char>(connectRequestRecord.address.size());
		connectResponse += connectRequestRecord.address;
		connectResponse += static_cast<unsigned char>(connectRequestRecord.port / 256);
		connectResponse += static_cast<unsigned char>(connectRequestRecord.port % 256);


		boost::asio::async_write(socket(),
			boost::asio::buffer(connectResponse.data(), connectResponse.size()),
			[this, self](boost::system::error_code ec, std::size_t length)
		{

			if (!ec)
			{
				transferDataForward();
				transferDataBackward();
			}
			else
			{
				lowerSocket().close();
				outsideConnectSocket.close();
			}
		});

	}


	void sendConnectResponseIpAddress(ConnectRequestRecord connectRequestRecord, std::array<unsigned char, 4> ipAddress)
	{
		auto self(shared_from_this());

		std::string connectResponse;

		connectResponse += 0x05;
		connectResponse += static_cast<char>(0x00);
		connectResponse += static_cast<char>(0x00);
		connectResponse += 0x01;

		connectResponse += ipAddress[0];
		connectResponse += ipAddress[1];
		connectResponse += ipAddress[2];
		connectResponse += ipAddress[3];
		connectResponse += static_cast<unsigned char>(connectRequestRecord.port / 256);
		connectResponse += static_cast<unsigned char>(connectRequestRecord.port % 256);


		boost::asio::async_write(socket(),
			boost::asio::buffer(connectResponse.data(), connectResponse.size()),
			[this, self](boost::system::error_code ec, std::size_t length)
		{

			if (!ec)
			{
				transferDataForward();
				transferDataBackward();
			}
			else
			{
				lowerSocket().close();
				outsideConnectSocket.close();
			}
		});

	}

	unsigned char forwardChar;
	unsigned char backwardChar;


	void transferDataForward()
	{
		auto self(shared_from_this());

		boost::asio::async_read(socket(),
			boost::asio::buffer(&forwardChar, 1),
			[this, self](boost::system::error_code ec, std::size_t /*length*/)
		{
			if (!ec)
			{
				boost::asio::async_write(outsideConnectSocket,
					boost::asio::buffer(&forwardChar, 1),
					[this, self](boost::system::error_code ec, std::size_t length)
				{

					if (!ec)
					{
						transferDataForward();
					}
					else
					{
						outsideConnectSocket.close();
						lowerSocket().close();
					}
				});
			}
			else
			{
				outsideConnectSocket.close();
				lowerSocket().close();
			}
		});
	}

	void transferDataBackward()
	{
		auto self(shared_from_this());

		boost::asio::async_read(outsideConnectSocket,
			boost::asio::buffer(&backwardChar, 1),
			[this, self](boost::system::error_code ec, std::size_t /*length*/)
		{
			if (!ec)
			{
				boost::asio::async_write(socket(),
					boost::asio::buffer(&backwardChar, 1),
					[this, self](boost::system::error_code ec, std::size_t length)
				{

					if (!ec)
					{
						transferDataBackward();
					}
					else
					{
						outsideConnectSocket.close();
						lowerSocket().close();
					}
				});
			}
			else
			{
				outsideConnectSocket.close();
				lowerSocket().close();
			}
		});
	}

	boost::asio::io_service& ioService;

	std::shared_ptr<ssl_socket> mSocket;

	boost::asio::ip::tcp::socket outsideConnectSocket;

	
};





class ProxyServer
{
public:
	ProxyServer(boost::asio::io_service& inIoService,
		const boost::asio::ip::tcp::endpoint& endpoint,
		boost::asio::ssl::context& sslContext)
		: ioService(inIoService)
		, acceptor(inIoService, endpoint)
		, socket(std::make_shared<ssl_socket>(inIoService, sslContext))
	{
		doAccept();
	}

private:
	void doAccept()
	{
		acceptor.async_accept(socket->lowest_layer(),
			[this](boost::system::error_code ec)
		{
			if (!ec)
			{
				std::make_shared<ProxySession>(socket, ioService)->start();
			}

			counter++;

			doAccept();
		});
	}

	boost::asio::io_service& ioService;

	boost::asio::ip::tcp::acceptor acceptor;
	std::shared_ptr<ssl_socket> socket;

	//std::map<size_t, ProxySession> proxySessionMap;

	size_t counter = 0;
};


int main()
{

	try
	{

		boost::asio::io_service ioService;

		boost::asio::ip::tcp::endpoint endpoint(boost::asio::ip::tcp::v4(), 8043);


		boost::asio::ssl::context sslContext(boost::asio::ssl::context::sslv23);

		sslContext.set_options(
			boost::asio::ssl::context::default_workarounds
			| boost::asio::ssl::context::no_sslv2
			| boost::asio::ssl::context::single_dh_use);


		std::function<std::string(std::size_t, boost::asio::ssl::context_base::password_purpose)> f = [](std::size_t, boost::asio::ssl::context_base::password_purpose) -> std::string { return ""; };
		sslContext.set_password_callback(f);

		sslContext.use_certificate_chain_file("server.crt");

		sslContext.use_private_key_file("server.key", boost::asio::ssl::context::pem);

		sslContext.use_tmp_dh_file("dh2048.pem");


		ProxyServer proxyServer(ioService, endpoint, sslContext);

		ioService.run();

	}
	catch (std::exception& e)
	{
		std::cerr << "Exception: " << e.what() << "\n";
	}


	return 0;

}