LeafWeb/Core/Remote/PiscalSshClient.cs

using System;
using System.Collections.Generic;
using System.Data.Common;
using System.IO;
using System.Linq;
using log4net;
using LeafWeb.Core.Utility;
using Polly;
using Renci.SshNet;
using Renci.SshNet.Common;

namespace LeafWeb.Core.Remote
{
	public class PiscalSshClient : IPiscalClient
	{
		private const string BaseDirectory = "./LeafWeb";
		private const string RemoteScriptPath = BaseDirectory + "/piscal_manager.sh";
		private readonly string _host;
		private readonly PasswordConnectionInfo _connectionInfo;

		private const string StatusComplete = "complete";
		private const string StatusRunning = "running";
		private const string StatusNotStarted = "not started";

		private readonly ILog _logger = LogManager.GetLogger(System.Reflection.MethodBase.GetCurrentMethod().DeclaringType);

        private readonly int _connectionRetryCount = 3;
        private readonly Policy _connectRetryPolicy;

		public PiscalSshClient(string connectionString)
		{
            if (string.IsNullOrEmpty(connectionString))
                throw new ArgumentNullException("connectionString");
			var conn = new DbConnectionStringBuilder {ConnectionString = connectionString};

			_host = conn["host"] as string;
			var username = conn["username"] as string;
			var password = conn["password"] as string;
			_connectionInfo = new PasswordConnectionInfo(_host, username, password);

            _connectRetryPolicy = Policy
                .Handle<SshConnectionException>()
                .Retry(_connectionRetryCount, 
                    (ex, i) => _logger.Warn($"Retry {i} after exception: {ex.Message}"));
        }

		private SshClient GetSshClient()
		{
			return new SshClient(_connectionInfo);
		}

		private ScpClient GetScpClient()
		{
			return new ScpClient(_connectionInfo);
		}
        
		private void Connect(PiscalLeafInput leafInput, BaseClient client)
		{
			try
            {
                _connectRetryPolicy.Execute(client.Connect);
            }
			catch (Exception ex)
			{
				var message = $"Exception while connecting to client. Message: {ex.Message}";
				_logger.Error(message);
				throw new PiscalClientException(leafInput.LeafInputId, message);
			}
		}

		private void Disconnect(PiscalLeafInput leafInput, BaseClient client)
		{
			try
			{
                _connectRetryPolicy.Execute(client.Disconnect);
			}
			catch (Exception ex)
			{
				var message = $"Exception while disconnecting from client. Message: {ex.Message}";
				_logger.Error(message);
				throw new PiscalClientException(leafInput.LeafInputId, message);
			}
		}

		private void CopyLeafInput(PiscalLeafInput leafInput)
		{
            // create directory
            using (var sshClient = GetSshClient())
            {
                Connect(leafInput, sshClient);
                sshClient.RunCommand($"mkdir {BaseDirectory}/{leafInput.PiscalDirectoryName}");
                sshClient.RunCommand($"mkdir {BaseDirectory}/{leafInput.PiscalDirectoryName}/input");
                Disconnect(leafInput, sshClient);
            }
            // copy files
			using (var scp = GetScpClient())
			foreach (var file in leafInput.InputFiles)
			{
				var inputPath = $"{BaseDirectory}/{leafInput.PiscalDirectoryName}/input/{file.Filename}";
				using (var stream = new MemoryStream(file.Contents))
				{
					_logger.Debug("Copying " + inputPath);
					Connect(leafInput, scp);
					scp.Upload(stream, inputPath);
					Disconnect(leafInput, scp);
				}
			}
		}

		public string Host => _host;

		public void RunLeafInput(PiscalLeafInput leafInput)
		{
			if (string.IsNullOrEmpty(leafInput.PhotosyntheticType))
				throw new PiscalClientException(leafInput.LeafInputId, "No PhotosyntheticType set");

			CopyLeafInput(leafInput);

			// begin processing
			using (var ssh = GetSshClient())
			{
				Connect(leafInput, ssh);
				var commandText = $"{RemoteScriptPath} -d {leafInput.PiscalDirectoryName} -p {leafInput.PhotosyntheticType} -s";

				if (leafInput.SuppressStorageCopy)
					commandText += " -t";

				if (!string.IsNullOrEmpty(leafInput.NotifyCompleteUrl))
					commandText += $" -u {leafInput.NotifyCompleteUrl}";

				var command = ssh.CreateCommand(commandText);
				command.Execute();
				Disconnect(leafInput, ssh);

				if (command.ExitStatus != 0)
					throw new PiscalClientException(leafInput.LeafInputId, command);

				var result = command.Result.TrimEndNewLine();
				if (result == "started")
					_logger.Info($"RunLeafInput result: {result}");
				else
					_logger.Warn($"RunLeafInput result: {result}, commandText: {commandText}");
			}
		}

		public PiscalStatus GetLeafInputStatus(PiscalLeafInput leafInput)
		{
			var statusRaw = GetLeafInputStatusRaw(leafInput);

			switch (statusRaw[0])
			{
				case StatusRunning:
					return PiscalStatus.Running;
				case StatusComplete:
					return PiscalStatus.Complete;
				case StatusNotStarted:
					return PiscalStatus.NotStarted;
				default:
					throw new PiscalClientException(leafInput.LeafInputId, "Unknown status: " + statusRaw[0]);
			}

		}

		private string[] GetLeafInputStatusRaw(PiscalLeafInput leafInput)
		{
			using (var ssh = GetSshClient())
			{
				Connect(leafInput, ssh);
				var commandText = $"{RemoteScriptPath} -d {leafInput.PiscalDirectoryName}";
				var command = ssh.CreateCommand(commandText);
				command.Execute();
				Disconnect(leafInput, ssh);

				if (command.ExitStatus != 0)
					throw new PiscalClientException(leafInput.LeafInputId, command);

				return command.Result
					.SplitNewLine()
					.Where(s => s.Length > 0)
					.Select(s => s.Trim()).ToArray();
			}
		}

		/// <summary>
		/// Gets the leaf output from piscal, only run on if result status is success
		/// </summary>
		public IEnumerable<PiscalLeafOutputFile> RetrieveLeafOutput(PiscalLeafInput leafInput)
		{
			// get output files
			var status = GetLeafInputStatusRaw(leafInput);
			if (status[0] != StatusComplete)
				throw new PiscalClientException(leafInput.LeafInputId, "Output not available, status is not StatusComplete, instead is currently " + status[0]);

			var filePaths = status.Skip(1);

			using (var scp = GetScpClient())
			{
				Connect(leafInput, scp);
				string outputFileType = string.Empty;
				foreach (var filePath in filePaths)
				{
					if (filePath.StartsWith("#"))
					{
						outputFileType = filePath.Substring(1);
						continue;
					}
					using (var stream = new MemoryStream())
					{
						scp.Download(filePath, stream);
						yield return 
							new PiscalLeafOutputFile
							{
								Contents = stream.ToArray(),
								Filename = filePath.FilenameFromPath(),
								PiscalDirectoryName = leafInput.PiscalDirectoryName,
								OutputFileType = outputFileType
							};
					}
				}

				Disconnect(leafInput, scp);
			}
		}

		public void KillLeafProcess(PiscalLeafInput leafInput)
		{
			using (var ssh = GetSshClient())
			{
				Connect(leafInput, ssh);
				var commandText = $"{RemoteScriptPath} -d {leafInput.PiscalDirectoryName} -k";
				var command = ssh.CreateCommand(commandText);
				command.Execute();
				Disconnect(leafInput, ssh);

				if (command.ExitStatus != 0)
					throw new PiscalClientException(leafInput.LeafInputId, command);
			}
		}

		public void CleanupLeafProcess(PiscalLeafInput leafInput)
		{
			using (var ssh = GetSshClient())
			{
				Connect(leafInput, ssh);
				var commandText = $"{RemoteScriptPath} -d {leafInput.PiscalDirectoryName} -c";
				var command = ssh.CreateCommand(commandText);
				command.Execute();
				Disconnect(leafInput, ssh);

				if (command.ExitStatus != 0)
					throw new PiscalClientException(leafInput.LeafInputId, command);
			}
		}
	}
}