CheckCircuit – Rev 2

Subversion Repositories:
Rev:
using System;
using System.Collections.Concurrent;
using System.IO;
using System.Net.Sockets;
using System.Net;
using System.Threading.Tasks;
using System.Threading;
using CheckCircuitCode.CommandLine;
using CommandLine;

namespace CheckCircuitCode
{
    internal class Program
    {
        private enum CIRCUIT_STATUS : int
        {
            NONE = 0,
            AUTHENTICATION_FAILED,
            CIRCUIT_UNAVAILABLE,
            COMMUNICATION_ERROR

        }

        public static TcpClient TcpClient { get; set; }
        public static bool Verbose { get; private set; }

        static int Main(string[] args)
        {
            var CancellationTokenSource = new CancellationTokenSource(TimeSpan.FromSeconds(60));
            var CancellationToken = CancellationTokenSource.Token;

            return Parser.Default.ParseArguments<ConnectOptions>(args)
                .MapResult(opts => RunConnectAndReturnExitCode(opts, CancellationToken), errs => 1);
        }

        private static int RunConnectAndReturnExitCode(ConnectOptions opts, CancellationToken cancellationToken)
        {
            Verbose = opts.Verbose;

            if (!IPEndPoint.TryParse(opts.Listen, out var listenEndPoint))
            {
                Console.WriteLine("Bad listen HOST:PORT format.");
                return 1;
            }

            if (!IPEndPoint.TryParse(opts.Connect, out var connectEndPoint))
            {
                Console.WriteLine("Bad connect HOST:PORT format.");
                return 1;
            }

            if (string.IsNullOrEmpty(opts.Password))
            {
                Console.WriteLine("Bad password string.");
                return 1;
            }

            if (!int.TryParse(opts.Timeout, out var timeout))
            {
                Console.WriteLine("Bad timeout value.");
                return 1;
            }

            var circuitStatus = Task.Run(() => StartClient(connectEndPoint, opts.Password, cancellationToken));

            return (int)circuitStatus.Result;
        }

        private static async Task<CIRCUIT_STATUS> StartClient(IPEndPoint connect, string password, CancellationToken cancellationToken)
        {
            try
            {
                using (var tcpClient = new TcpClient())
                {
                    await tcpClient.ConnectAsync(connect.Address, connect.Port, cancellationToken);
                    using (var tcpClientStream = tcpClient.GetStream())
                    {
                        using (var streamReader = new StreamReader(tcpClientStream))
                        {
                            using (var streamWriter = new StreamWriter(tcpClientStream))
                            {
                                do
                                {
                                    if (!await Authenticate(streamReader, streamWriter, password))
                                    {
                                        if (Verbose)
                                        {
                                            Console.WriteLine("Authentication failed, please check OR port and password.");
                                        }
                                        return CIRCUIT_STATUS.AUTHENTICATION_FAILED;
                                    }

                                    do
                                    {
                                        try
                                        {
                                            if (await IsCircuitFormed(streamReader, streamWriter))
                                            {
                                                return CIRCUIT_STATUS.NONE;
                                            }

                                            return CIRCUIT_STATUS.CIRCUIT_UNAVAILABLE;
                                        }
                                        catch (ArgumentException)
                                        {
                                            return CIRCUIT_STATUS.COMMUNICATION_ERROR;
                                        }

                                    } while (tcpClient.Connected && !cancellationToken.IsCancellationRequested);
                                } while (tcpClient.Connected && !cancellationToken.IsCancellationRequested);
                            }
                        }
                    }
                }
            }
            catch (Exception exception)
            {
                if (Verbose)
                {
                    Console.WriteLine(exception);
                }

                return CIRCUIT_STATUS.COMMUNICATION_ERROR;
            }
        }

        public static async Task<bool> Authenticate(StreamReader sr, StreamWriter sw, string password)
        {
            var readLineTask = sr.ReadLineAsync();
            await sw.WriteLineAsync($"AUTHENTICATE \"{password}\"");
            await sw.FlushAsync();
            var readLine = await readLineTask;

            return string.Equals(readLine, @"250 OK", StringComparison.Ordinal);
        }

        public static async Task<bool> IsCircuitFormed(StreamReader sr, StreamWriter sw)
        {
            var readLineTask = sr.ReadLineAsync();
            await sw.WriteLineAsync(@"GETINFO status/circuit-established");
            await sw.FlushAsync();
            var readLine = await readLineTask;

            var success = string.Equals(readLine, @"250-status/circuit-established=1", StringComparison.Ordinal);
            readLine = await sr.ReadLineAsync();
            if (string.Equals(readLine, @"250 OK", StringComparison.Ordinal))
            {
                return success;
            }

            throw new ArgumentException("Unable to read response from control port.");
        }

        public static async Task SendPayload(ConcurrentDictionary<IntPtr, TcpClient> clients, string payload, CancellationToken CancellationToken)
        {
            foreach (var (handle, client) in clients)
            {
                if (CancellationToken.IsCancellationRequested)
                {
                    throw new TaskCanceledException();
                }

                try
                {
                    if (!client.Connected)
                    {
                        continue;
                    }

                    using (var networkStream = client.GetStream())
                    {
                        using (var sw = new StreamWriter(networkStream))
                        {
                            await sw.WriteLineAsync(payload);

                        }
                    }
                }
                catch (Exception exception)
                {
                    if (Verbose)
                    {
                        Console.WriteLine(exception);
                    }
                }
                finally
                {
                    clients.TryRemove(handle, out _);
                    try
                    {
                        client.Close();
                    }
                    catch (ObjectDisposedException exception)
                    {
                        if (Verbose)
                        {
                            Console.WriteLine(exception);
                        }
                    }
                }
            }
        }
    }
}