using System;
using System.Diagnostics;
using System.Linq;
using System.Net;
using System.Net.Sockets;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace TcpTunnel
{
    internal class Program
    {
        private static void Main(string[] args)
        {
            //#if DEBUG
            //            Debugger.Launch();
            //            Debugger.Break();
            //#endif

            if (args.Length != 6)
            {
                Console.WriteLine($"Usage: <cl|sv> <token> <listenAddress> <listenPort> <targetHost> <targetPort>");
                return;
            }

            var isCl = args[0] == "cl";
            var token = args[1];

            if (!IPAddress.TryParse(args[2], out var listenAddress))
            {
                Console.WriteLine("invalid listen address");
                return;
            }

            if (!int.TryParse(args[3], out var listenPort))
            {
                Console.WriteLine("invalid listen port");
                return;
            }

            var targetHost = args[4];

            if (!int.TryParse(args[5], out var targetPort))
            {
                Console.WriteLine("invalid target port");
                return;
            }

            var listener = new TcpListener(listenAddress, listenPort);
            listener.Start(1);

            Console.WriteLine($"Listening on {listener.LocalEndpoint}");
            var connectionCount = 0;

            var tokenBuf = Encoding.UTF8.GetBytes(token);

            listener.BeginAcceptTcpClient(AsyncCallback, null);
            void AsyncCallback(IAsyncResult ar)
            {
                var cn = Interlocked.Increment(ref connectionCount);


                listener.BeginAcceptTcpClient(AsyncCallback, null);

                var ibClient = listener.EndAcceptTcpClient(ar);


                Console.WriteLine($"[#{cn:0000}]Income from {ibClient.Client.RemoteEndPoint}");
                var ibStream = ibClient.GetStream();

                NetworkStream obStream = null;
                TcpClient obClient = null;
                if (isCl)
                {
                    Console.WriteLine($"[#{cn:0000}]Connecting to target...");
                    obClient = new TcpClient(targetHost, targetPort);
                    obStream = obClient.GetStream();

                    Console.WriteLine($"[#{cn:0000}]Connected sending token...");
                    obStream.WriteByte((byte)tokenBuf.Length);
                    obStream.Write(tokenBuf, 0, tokenBuf.Length);
                }
                else
                {
                    Console.WriteLine($"[#{cn:0000}]Checking token...");

                    var tol = ibStream.ReadByte();
                    if (tol != tokenBuf.Length)
                    {
                        Console.WriteLine($"[#{cn:0000}]Invalid token length closing");

                        Thread.Sleep(1000);
                        ibStream.Close();
                        ibClient.Close();

                        Console.WriteLine($"[#{cn:0000}]Invalid token length closed");
                        return;
                    }

                    var buf = new byte[tol];
                    ibStream.Read(buf, 0, tol);
                    if (!buf.SequenceEqual(tokenBuf))
                    {
                        Console.WriteLine($"[#{cn:0000}]Invalid token closing");

                        Thread.Sleep(1000);
                        ibStream.Close();
                        ibClient.Close();

                        Console.WriteLine($"[#{cn:0000}]Invalid token closed");
                        return;
                    }

                    Console.WriteLine($"[#{cn:0000}]Auth OK, Connecting to target...");
                    obClient = new TcpClient(targetHost, targetPort);
                    obStream = obClient.GetStream();
                }

                if (ibClient.Connected && obClient.Connected)
                {
                    Console.WriteLine($"[#{cn:0000}]Start tunneling");
                    
                    ibClient.ReceiveTimeout = 120000;
                    ibClient.SendTimeout = 120000;
                    obClient.ReceiveTimeout = 120000;
                    obClient.SendTimeout = 120000;

                    var i2o = Task.Factory.StartNew(delegate { ibStream.CopyTo(obStream); });
                    var o2i = Task.Factory.StartNew(delegate { obStream.CopyTo(ibStream); });

                    Task.WaitAny(i2o, o2i);
                    Console.WriteLine($"[#{cn:0000}]End tunnel");
                }
                else
                {
                    Console.WriteLine($"[#{cn:0000}]No connected.");
                }

                ibClient?.Dispose();
                obClient?.Dispose();
            }

            Console.WriteLine("Press ENTER to Stop");
            Console.ReadLine();
            Console.WriteLine("Stopping...");
            listener.Stop();
            Console.WriteLine("Stopped,press ENTER to exit");
            Console.ReadLine();
        }
    }
}