path: root/Software/Visual_Studio/Utilities/Tango.RemoteRunner.UI/MainWindowVM.cs

using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Diagnostics;
using System.IO;
using System.IO.Compression;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using System.Windows;
using Tango.Logging;
using Tango.PMR;
using Tango.PMR.Common;
using Tango.PMR.Discovery;
using Tango.PMR.IO;
using Tango.SharedUI;
using Tango.Transport;
using Tango.Transport.Adapters;
using Tango.Transport.Discovery;
using Tango.Transport.Servers;
using Tango.Transport.Transporters;

namespace Tango.RemoteRunner.UI
{
    public class MainWindowVM : ViewModel
    {
        private TcpServer _server;
        private ITransporter _transporter;
        private UdpDiscoveryService<BasicDiscoveryMessage> _discoveryService;
        private const int PORT = 9595;
        private List<FileUpload> _uploads;
        private const long MAX_CHUNK_LENGTH = 500000; //500kb

        private Action<String> _notificationAction;
        private ObservableCollection<RunningProcess> _runningProcesses;
        public ObservableCollection<RunningProcess> RunningProcesses
        {
            get { return _runningProcesses; }
            set { _runningProcesses = value; RaisePropertyChangedAuto(); }
        }

        private String _log;
        public String Log
        {
            get { return _log; }
            set { _log = value; RaisePropertyChangedAuto(); }
        }

        public MainWindowVM()
        {
            RunningProcesses = new ObservableCollection<RunningProcess>();
            _uploads = new List<FileUpload>();

            var logger = new SimpleStringLogger() { Enabled = true };
            logger.LogReceived += (_, log) =>
            {
                Log += String.Format("{0}: {1}", log.TimeStamp.ToShortTimeString(), log.Message) + Environment.NewLine;
            };
            LogManager.RegisterLogger(logger);
        }

        public void ViewLoaded()
        {
            LogManager.Log("Application Started!");
            LogManager.Log("Initializing TCP listener on port " + PORT + "...");
            _server = new TcpServer(PORT);
            _server.ClientConnected += _server_ClientConnected;
            _server.Start();

            LogManager.Log("Initializing basic transporter...");
            _transporter = new BasicTransporter();
            _transporter.FailsWithAdapter = false;
            _transporter.RequestReceived += _transporter_RequestReceived;

            _discoveryService = new UdpDiscoveryService<BasicDiscoveryMessage>(2018, new BasicDiscoveryMessage()
            {
                ServiceName = "Tango Remote Runner",
                Port = PORT,
            });

            _discoveryService.Interval = TimeSpan.FromSeconds(1);
            _discoveryService.Start();
        }

        private void _transporter_RequestReceived(object sender, MessageContainer container)
        {
            try
            {
                switch (container.Type)
                {
                    case MessageType.FileUploadRequest:
                        HandleFileUploadRequest(MessageFactory.ParseTangoMessageFromContainer<FileUploadRequest>(container));
                        break;
                    case MessageType.FileChunkUploadRequest:
                        HandleFileChunkUploadRequest(MessageFactory.ParseTangoMessageFromContainer<FileChunkUploadRequest>(container));
                        break;
                    case MessageType.ExecuteProcessRequest:
                        HandleExecuteProcessRequest(MessageFactory.ParseTangoMessageFromContainer<ExecuteProcessRequest>(container));
                        break;
                    case MessageType.KillProcessRequest:
                        HandleKillProcessRequest(MessageFactory.ParseTangoMessageFromContainer<KillProcessRequest>(container));
                        break;
                }
            }
            catch (Exception ex)
            {
                LogManager.Log(ex);
                _transporter.SendErrorResponse(ex, container.Token);
            }
        }

        #region Request Handlers

        private void HandleFileUploadRequest(TangoMessage<FileUploadRequest> request)
        {
            LogManager.Log("File upload request received " + request.Message.Path);

            Notify("Upload request " + request.Message.Path);

            var tempFolder = TemporaryManager.CreateFolder();
            LogManager.Log("Created temporary upload folder " + tempFolder.Path);

            String uploadID = Guid.NewGuid().ToString();

            _uploads.Add(new FileUpload()
            {
                Folder = tempFolder,
                UploadID = uploadID,
                FileName = request.Message.Path,
                Length = request.Message.Length,
            });

            _transporter.SendResponse<FileUploadResponse>(new FileUploadResponse()
            {
                UploadID = uploadID,
                MaxChunkLength = MAX_CHUNK_LENGTH,
            }, request.Container.Token);
        }

        private void HandleFileChunkUploadRequest(TangoMessage<FileChunkUploadRequest> request)
        {
            FileUpload upload = _uploads.SingleOrDefault(x => x.UploadID == request.Message.UploadID);
            if (upload == null)
            {
                _transporter.SendResponse<FileChunkUploadResponse>(new FileChunkUploadResponse() { IsCanceled = true }, request.Container.Token, true, ErrorCode.InvalidUploadId);
                return;
            }

            if (upload.CurrentLength + request.Message.Buffer.Length > upload.Length)
            {
                _transporter.SendResponse<FileChunkUploadResponse>(new FileChunkUploadResponse() { IsCanceled = true }, request.Container.Token, true, ErrorCode.FileLengthOutOfRange);
                return;
            }

            using (FileStream fs = new FileStream(upload.FullPath, FileMode.Append))
            {
                byte[] buffer = request.Message.Buffer.ToByteArray();
                fs.Write(buffer, 0, buffer.Length);
                upload.CurrentLength = fs.Length;

                if (fs.Length == upload.Length)
                {
                    upload.Completed = true;
                    LogManager.Log("File upload completed " + upload.FileName);
                }
            }

            _transporter.SendResponse<FileChunkUploadResponse>(new FileChunkUploadResponse(), request.Container.Token);
        }

        private void HandleExecuteProcessRequest(TangoMessage<ExecuteProcessRequest> request)
        {
            LogManager.Log("Execute process request received " + request.Message.FileName);

            Notify("Execute request " + request.Message.FileName);

            FileUpload upload = _uploads.SingleOrDefault(x => x.UploadID == request.Message.UploadID);
            if (upload == null || !upload.Completed)
            {
                _transporter.SendResponse<ExecuteProcessResponse>(new ExecuteProcessResponse(), request.Container.Token, true, ErrorCode.InvalidUploadId);
                return;
            }

            if (Path.GetExtension(upload.FileName).ToLower() == ".zip" && !upload.Expanded)
            {
                LogManager.Log("Extracting zip file to " + upload.Folder.Path);
                ZipFile.ExtractToDirectory(upload.FullPath, upload.Folder.Path);
            }

            upload.Expanded = true;

            RunningProcess process = new RunningProcess();
            process.Process = Process.Start(Path.Combine(upload.Folder.Path, Path.GetFileName(request.Message.FileName)));

            process.Process.Exited += (_, __) => 
            {
                RunningProcesses.Remove(process);
            };

            LogManager.Log("Process started " + process.ProcessName + ", " + process.ProcessID);

            InvokeUINow(() =>
            {
                RunningProcesses.Add(process);
            });

            _transporter.SendResponse<ExecuteProcessResponse>(new ExecuteProcessResponse() { ProcessID = process.ProcessID }, request.Container.Token);
        }

        private void HandleKillProcessRequest(TangoMessage<KillProcessRequest> request)
        {
            LogManager.Log("Kill process request received " + request.Message.ProcessID);
            Notify("Kill request " + request.Message.ProcessID);

            var process = RunningProcesses.SingleOrDefault(x => x.ProcessID == request.Message.ProcessID);

            if (process == null)
            {
                _transporter.SendResponse<KillProcessResponse>(new KillProcessResponse(), request.Container.Token, true, ErrorCode.InvalidProcessId);
                return;
            }

            process.Process.Kill();
            LogManager.Log("Process killed " + process.ProcessName + ", " + process.ProcessID);
            _transporter.SendResponse<KillProcessResponse>(new KillProcessResponse(), request.Container.Token);
        }

        #endregion

        private void _server_ClientConnected(object sender, ClientConnectedEventArgs e)
        {
            Task.Factory.StartNew(async () =>
            {
                try
                {
                    Stopwatch watch = new Stopwatch();
                    watch.Start();

                    LogManager.Log("Remote client connection from " + e.Socket.GetIPAddress());
                    LogManager.Log("Initializing transporter with a new TCP adapter...");

                    Notify("Client connected from " + e.Socket.GetIPAddress());

                    await _transporter.Disconnect();
                    TcpTransportAdapter adapter = new TcpTransportAdapter(e.Socket);
                    _transporter.Adapter = adapter;
                    await _transporter.Connect();

                    LogManager.Log(watch.Elapsed.Milliseconds + " milliseconds");
                }
                catch (Exception ex)
                {
                    LogManager.Log(ex);
                }
            });
        }

        public void SetNotificationAction(Action<String> notificationAction)
        {
            _notificationAction = notificationAction;
        }

        private void Notify(String text)
        {
            _notificationAction?.Invoke(text);
        }
    }
}