///
/// Copyright (c) 2017 Sensus Metering Systems
///
using System;
using System.Net;
using System.Net.Sockets;
using System.Threading;
using System.Collections.Generic;
using System.Text;
using System.IO;
using System.Diagnostics;
using log4net;

namespace TBF.BenchControl.Network.Tftp
{
	/// <summary>
	/// Simple TFTP Server	
	/// </summary>		
	public class TftpServer
	{
		static readonly ILog log = LogManager.GetLogger("TftpServer");

		const ushort DefaultTftpPortNr = 69;
		const int TftpHeaderLength = 4;
		const int TftpDataBlockLength = 512;
		readonly char[] Delimiter = new char[] { '\0' };		/// Delimits filename, transfer mode and error message in UDP packets

		/// <summary>
		/// TFTP server request modes
		/// </summary>
		enum RequestMode
		{
			BinaryRead,
			BinaryWrite,
			AsciiRead,
			AsciiWrite,
		}

		/// <summary>
		/// Opcodes identifying TFTP packets
		/// </summary>
		enum Opcode : short
		{
			RRQ		= 01,		/// Read request packet.
			WRQ		= 02,		/// Write request packet. 
			DATA	= 03,		/// Data packet. 
			ACK		= 04,		/// Acknowledgement packet.
			ERROR	= 05,		/// Wrror packet. 
			OACK	= 06		/// Option acknowledgement packet. 
		}

		/// Strings included in RRQ/WRQ packets
		const string REQUEST_MODE_NETASCII = "netascii";
		const string REQUEST_MODE_BINARY = "octet";

		/// <summary>
		/// Error codes included in ERROR packet
		/// </summary>
		enum ErrorCode : short
		{
			NO_ERROR				= 0,	/// Not defined, see error message (if any).
			ERROR_FILE_NOT_FOUND	= 1,	/// File not found.
			ERROR_ACCESS_VIOLATION	= 2,	/// Access violation.
			ERROR_ALLOC_ERROR		= 3,	/// Disk full or allocation exceeded.
			ERROR_ILLEGAL_OP		= 4,	/// Illegal TFTP operation.
			ERROR_UNKNOWN_TID		= 5,	/// Unknown transfer ID.
			ERROR_FILE_EXISTS		= 6,	/// File already exists.
			ERROR_INVALID_USER		= 7,	/// No such user.
		}


		///
		/// Single file transfer related class
		///
		class FileTransfer
		{
			public IPEndPoint Endpoint { get { return endpoint; } }
			IPEndPoint endpoint;

			public string Filename { get { return filename; } }
			string filename;

			public RequestMode Mode { get { return mode; } }
			RequestMode mode;

			public System.IO.Stream Stream;
			public BinaryReader StreamReader;
			public BinaryWriter StreamWriter;
			public int CurrentBlock;
			public int LastSentDataLength;

			public FileTransfer(IPEndPoint endpoint, string filename, RequestMode mode)
			{
				this.endpoint = endpoint;
				this.filename = filename;
				this.mode = mode;

				Stream = null;
				StreamReader = null;
				StreamWriter = null;
				CurrentBlock = (mode == RequestMode.BinaryRead ? 1 : 0);
				LastSentDataLength = 0;
			}

			/// <summary>
			/// IPEndPoint can be used as a key as there is max. one file transfer for each endpoint.
			/// </summary>
			/// <param name="ep">IP endpoint</param>
			/// <param name="list">List of FileTransfer-s</param>
			/// <returns></returns>
			public static FileTransfer Find(IPEndPoint ep, IList<FileTransfer> list)
			{
				if (list == null) return null;
				foreach (FileTransfer ft in list)
				{
					if ((ft.endpoint.Address.Equals(ep.Address)) && (ft.endpoint.Port.Equals(ep.Port)))
					{
						return ft;
					}
				}
				return null;
			}

			/// <summary>
			/// FileTransfer Close() closes open streams.
			/// </summary>
			public void Close()
			{
				if (StreamReader != null) StreamReader.Close();
				if (StreamWriter != null) StreamWriter.Close();
				if (Stream != null) Stream.Close();
			}

			/// For debugging.
			public override string ToString()
			{
				return	"[" + endpoint.ToString() + "," + filename + "," + mode.ToString() +
						",block=" + CurrentBlock.ToString() + ",last=" + LastSentDataLength.ToString() + "]";
			}
		}


		///
		/// TFTP server instance related instance fields
		///
		string tftpDirectory;					/// TFTP root directory
		int listenPort;							/// UDP port to listen to
		UdpClient udpClient;
		Thread listenThread;
		bool done;								/// Flag to finish listenThread											
		IList<FileTransfer> fileTransfers;		/// Current file transfers (max. one transfer for each endpoint)


		public TftpServer()
		{
			tftpDirectory = null;
			udpClient = null;				/// UDP client
			listenThread = null;			/// Listen thread
			done = false;					/// Flag to finish listenThread
			fileTransfers = null;			/// No file transfer
		}


		#region Event handler

		public class NotificationEventArgs : EventArgs
		{
			public string Message;
			public IPAddress IP;
			public string FileName;

			public NotificationEventArgs(string message, IPAddress ip, string filename)
			{
				this.Message = message;
				this.IP = ip;
				this.FileName = filename;
			}
		}

		public delegate void NotificationEventHandler(object sender, NotificationEventArgs args);

		public static event NotificationEventHandler NotificationHandler;

		/// <summary>
		/// Writes a log and invokes notification handlers (updates TFTP tab-page).
		/// </summary>
		/// <param name="message">Notification text start</param>
		/// <param name="fileTransfer">File transfer object reference or null</param>
		/// <param name="logType">Logger log type </param>
		void OnNotify(string message, FileTransfer fileTransfer, log4net.Core.Level level)
		{
			/// Write a log
			IPAddress ip = null;
			string filename = null;
			string logStr = message;

			if (fileTransfer != null)
			{
				ip = fileTransfer.Endpoint.Address;
				filename = fileTransfer.Filename;
				logStr = string.Format("{0} IP={1}, filename={2}", message, ip, filename);
			}

			if (level == log4net.Core.Level.Debug)
				log.Debug(logStr);
			else if (level == log4net.Core.Level.Info)
				log.Info(logStr);
			else if (level == log4net.Core.Level.Warn)
				log.Warn(logStr);
			else if (level == log4net.Core.Level.Error)
				log.Error(logStr);
			else if (level == log4net.Core.Level.Fatal || level == log4net.Core.Level.All)
				log.Fatal(logStr);

			/// Notification: invoke handlers
			if (NotificationHandler != null)
			{
				NotificationHandler(null, new NotificationEventArgs(message, ip, filename));
			}
		}

		#endregion


		/// <summary>
		/// Start TFTP server
		/// </summary>
		/// <param name="tftpDirectory">TFTP root directory</param>
		/// <param name="netAdapterInfo">Network adadpter to listen to, null=listen to all net.adapters</param>
		/// <returns>true = TFTP is running, false = TFTP start failed</returns>
		public bool Start(string tftpDirectory, IPAddress ipAddress)
		{
			if (!Directory.Exists(tftpDirectory)) return false;	/// TFTP root directory does not exist								
			if (!tftpDirectory.EndsWith(Path.DirectorySeparatorChar.ToString()))
			{
				tftpDirectory += Path.DirectorySeparatorChar.ToString();
			}
			this.tftpDirectory = tftpDirectory;
			this.listenPort = DefaultTftpPortNr;

			if (udpClient != null) return false;		/// TFTP server is already running
			try
			{
				if (ipAddress != null)
				{
					IPEndPoint listenEndpoint = new IPEndPoint(ipAddress, listenPort);
					udpClient = new UdpClient(listenEndpoint);
				}
				else
				{
					udpClient = new UdpClient(listenPort, AddressFamily.InterNetwork);
				}
			}
			catch
			{
				return false;							/// Cannot open UDP port (is another TFTP running?)
			}

			try
			{
				fileTransfers = new List<FileTransfer>();
				listenThread = new Thread(new ThreadStart(Listener));
				listenThread.Name = "TftpServer";				/// Is used as a log file root name
				listenThread.Start();
			}
			catch
			{
				return false;
			}
			return true;
		}

		/// <summary>
		/// Stop TFTP server
		/// </summary>
		public void Stop()
		{
			done = true;

			if (udpClient != null)
			{
				udpClient.Close();
				udpClient = null;
			}

			try
			{
				/// Close the open file streams
				if (fileTransfers != null)
				{
					foreach (FileTransfer ft in fileTransfers) ft.Close();
				}
			}
			catch (Exception ex)
			{
				log.FatalFormat("Fatal error in Stop(): {0}", ex.Message);
			}
			fileTransfers = null;
		}


		#region Listner thread section

		/// <summary>
		/// Get readable string describing the packet.
		/// </summary>
		/// <param name="data">TFTP packet</param>
		/// <param name="actualLength">Number of valid byted in array data</param>
		/// <returns>description</returns>
		string GetDebugMessage(Byte[] data, int actualLength)
		{
			if (data == null || data.Length < actualLength || actualLength < TftpHeaderLength) return "[invalid packet]";

			Opcode opcode = (Opcode)((((short)data[0]) * 256) + (short)data[1]);

			if (opcode == Opcode.RRQ || opcode == Opcode.WRQ)
			{
				Encoding ASCII = Encoding.ASCII;
				string[] strData = ASCII.GetString(data, 2, data.Length - 2).Split(Delimiter, 3);
				string filename = strData[0];
				string mode = strData[1].ToLower();

				return "[" + opcode.ToString() + " filename=" + filename + ", mode=" + mode + "]";
			}

			int number = data[2] * 256 + data[3];

			if (actualLength > TftpHeaderLength)
			{
				return "[" + opcode.ToString() + " + " + number.ToString() + " + " + (actualLength - TftpHeaderLength).ToString() + " bytes]";
			}
			else
			{
				return "[" + opcode.ToString() + " + " + number.ToString() + "]";
			}
		}


		/// <summary>
		/// Main listener thread loop
		/// </summary>
		public void Listener() 
		{
			OnNotify("----- TFTP server started -----", null, log4net.Core.Level.Fatal);

			while (!done) 
			{
				try
				{
					IPEndPoint endpoint = null;
					Byte[] data = udpClient.Receive(ref endpoint);

					log.InfoFormat("Rcvd packet {0}", GetDebugMessage(data, data.Length));

					///
					/// Process the packet
					///
					Opcode opcode = (Opcode)((((short)data[0]) * 256) + (short)data[1]);

					///
					/// Find any existing file transfer with this IP endpoint
					///
					FileTransfer fileTransfer = FileTransfer.Find(endpoint, fileTransfers);

					if (opcode == Opcode.RRQ || opcode == Opcode.WRQ)
					{
						if (fileTransfer != null)
						{
							///
							/// If there is an incomplete file transfer with this endpoint, close the open streams.
							///
							OnNotify("Forcing an incomplete file transfer to close: ", fileTransfer, log4net.Core.Level.Error);

							fileTransfer.Close();
							fileTransfers.Remove(fileTransfer);
							fileTransfer = null;
						}
					}
					else if (fileTransfer == null)
					{
						/// If there is no existing file transfer with this endpoint, ignore packets other then RRQ or WRQ.
						log.ErrorFormat("No matching file transfer - packet ignored: IP={0}", endpoint.Address);
						continue;
					}


					switch (opcode)
					{
						case Opcode.RRQ:
							ProcessReadRequest(data, endpoint);
							break;

						case Opcode.WRQ:
							ProcessWriteRequest(data, endpoint);
							break;

						case Opcode.ERROR:
							ProcessError(data, fileTransfer);
							break;

						case Opcode.ACK:
							ProcessAck(data, fileTransfer);							
							break;

						case Opcode.DATA:
							ProcessData(data, fileTransfer);
							break;

						case Opcode.OACK:
						default:
							break;

					}	
				}
				catch (Exception exc)
				{
					log.FatalFormat("Fatal error in Listener(): {0}", exc.Message);
				}
			}
			return;
		}


		#region TFTP 'GET' section

		/// <summary>
		/// Handle 'Read Request' (RRQ)
		/// </summary>
		/// <param name="data">Data from RRQ packet</param>
		/// <param name="endpoint">Client IP</param>
		private void ProcessReadRequest(Byte[] data, IPEndPoint endpoint)
		{
			/// Extract a filename and a mode
			string[] strData = Encoding.ASCII.GetString(data, 2, data.Length - 2).Split(Delimiter, 3);
			string filename = strData[0];
			string mode = strData[1].ToLower();

			if (mode == REQUEST_MODE_BINARY)
			{
				/// Create FileTransfer object, do not add it to this.fileTransfers yet
				FileTransfer fileTransfer = new FileTransfer(endpoint, filename, RequestMode.BinaryRead);

				/// Try to open the source file
				try
				{
					fileTransfer.Stream = System.IO.File.OpenRead(tftpDirectory + filename);
					fileTransfer.StreamReader = new BinaryReader(fileTransfer.Stream);
				}
				catch
				{
					OnNotify("Cannot open file to be uploaded: ", fileTransfer, log4net.Core.Level.Warn);
					SendError(fileTransfer, ErrorCode.ERROR_FILE_NOT_FOUND);
					return;
				}

				/// File open OK, start the transfer, send the 1st data packet
				OnNotify("Uploading file:        ", fileTransfer, log4net.Core.Level.Info);
				fileTransfers.Add(fileTransfer);
				SendData(fileTransfer);
			}
			else
			{
				throw (new Exception("Built-in TFTP server does not support '" + mode + "' file transfer mode."));
			}
		}

		/// <summary>
		/// Send part of file data
		/// </summary>
		/// <param name="endpoint">location to send stream to</param>
		/// <param name="blockNumber">512 byte block to send</param>
		private void SendData(FileTransfer fileTransfer)
		{
			int fileOffset = (fileTransfer.CurrentBlock - 1) * TftpDataBlockLength;
			fileTransfer.Stream.Seek(fileOffset, SeekOrigin.Begin);

			/// Prepare data buffer
			Byte[] buffer = new Byte[TftpDataBlockLength + TftpHeaderLength];
			buffer[0] = 0;
			buffer[1] = (byte)Opcode.DATA;
			buffer[2] = (byte)((fileTransfer.CurrentBlock & 0x0000FF00) >> 8);
			buffer[3] = (byte) (fileTransfer.CurrentBlock & 0x000000FF);
			fileTransfer.LastSentDataLength = fileTransfer.StreamReader.Read(buffer, TftpHeaderLength, TftpDataBlockLength);

			/// Send the data packet
			int ecode = udpClient.Send(buffer, fileTransfer.LastSentDataLength + TftpHeaderLength, fileTransfer.Endpoint);
			log.InfoFormat("Sent packet {0} ecode={1}", GetDebugMessage(buffer, fileTransfer.LastSentDataLength + TftpHeaderLength), ecode);

			/// Check the return value
			if (ecode != fileTransfer.LastSentDataLength + TftpHeaderLength)
			{
				OnNotify("Error when sending data: ecode = " + ecode.ToString() + ", ", fileTransfer, log4net.Core.Level.Error);
				fileTransfer.Close();
				fileTransfers.Remove(fileTransfer);
			}
		}

		/// <summary>
		/// Handle ACK response and send next block.
		/// </summary>
		/// <param name="data">data from packet</param>
		/// <param name="endpoint">client</param>
		private void ProcessAck(Byte[] data, FileTransfer fileTransfer)
		{
			int protocolBlocknum = 256 * (int)data[2] + (int)data[3];

			/// Check the received block number, support for file size 2GB
			if (protocolBlocknum == (fileTransfer.CurrentBlock & 0xFFFF))
			{
				/// OK => Check the last data block length
				if (fileTransfer.LastSentDataLength < TftpDataBlockLength)
				{
					/// Nothing more to send, finish the transfer
					OnNotify("Upload completed: ", fileTransfer, log4net.Core.Level.Info);
					fileTransfer.Close();
					fileTransfers.Remove(fileTransfer);
				}
				else
				{
					/// Send the next data block
					fileTransfer.CurrentBlock++;
					SendData(fileTransfer);
				}
			}
			else
			{
				/// Block number NOK => re-send the last block (without incrementing fileTransfer.CurrentBlock).
				/// This often happens when transferring large files due to iCOM performance limits.
				log.ErrorFormat("Received ACK does not match the sent block#: IP={0} ---> Re-sending", fileTransfer.Endpoint.Address);
				SendData(fileTransfer);
			}
		}

		#endregion


		#region TFTP 'PUT' section

		/// <summary>
		/// Handle 'Write Request' (WRQ)
		/// </summary>
		/// <param name="data">Data from WRQ packet</param>
		/// <param name="endpoint">Client IP</param>
		private void ProcessWriteRequest(Byte[] data, IPEndPoint endpoint)
		{
			/// Extract a filename and a mode
			string[] strData = Encoding.ASCII.GetString(data, 2, data.Length - 2).Split(Delimiter, 3);
			string filename = strData[0];
			string mode = strData[1].ToLower();

			if (mode == REQUEST_MODE_BINARY)
			{
				/// Create FileTransfer object, do not add it to this.fileTransfers yet
				FileTransfer fileTransfer = new FileTransfer(endpoint, filename, RequestMode.BinaryWrite);

				/// Try to open the source file
				try
				{
					fileTransfer.Stream = System.IO.File.Create(tftpDirectory + filename);
					fileTransfer.StreamWriter = new BinaryWriter(fileTransfer.Stream);
				}
				catch
				{
					OnNotify("Cannot open file to be downloaded for writing: ", fileTransfer, log4net.Core.Level.Warn);
					SendError(fileTransfer, ErrorCode.ERROR_FILE_EXISTS);
					return;
				}

				/// File open OK, start the transfer, send the 0th acknowledge packet
				OnNotify("Downloading file:        ", fileTransfer, log4net.Core.Level.Info);
				fileTransfers.Add(fileTransfer);
				SendAck(fileTransfer);
			}
			else
			{
				throw (new Exception("Built-in TFTP server does not support '" + mode + "' file transfer mode."));
			}
		}

		/// <summary>
		/// Send acknowledge packet (ACK)
		/// </summary>
		/// <param name="fileTransfer">Current FileTransfer object</param>
		private void SendAck(FileTransfer fileTransfer)
		{
			const int AcknowledgePacketLength = 4;

			/// Prepare the acknowledgment packet
			Byte[] buffer = new Byte[AcknowledgePacketLength];
			buffer[0] = 0;
			buffer[1] = (byte)Opcode.ACK;
			buffer[2] = (byte)((fileTransfer.CurrentBlock & 0x0000FF00) >> 8);
			buffer[3] = (byte) (fileTransfer.CurrentBlock & 0x000000FF);

			/// Send the packet
			int ecode = udpClient.Send(buffer, AcknowledgePacketLength, fileTransfer.Endpoint);
			log.InfoFormat("Sent packet {0} ecode={1}", GetDebugMessage(buffer, AcknowledgePacketLength), ecode);

			/// Check the return value
			if (ecode != AcknowledgePacketLength)
			{
				OnNotify("Error when sending ACK packet: ecode = " + ecode.ToString() + ", ", fileTransfer, log4net.Core.Level.Error);
				fileTransfer.Close();
				fileTransfers.Remove(fileTransfer);
			}
		}

		private void ProcessData(Byte[] data, FileTransfer fileTransfer)
		{
			int protocolBlocknum = 256 * (int)data[2] + (int)data[3];

			/// Check the received data block number, support max. file size 2GB
			if (protocolBlocknum != ((fileTransfer.CurrentBlock + 1) & 0xFFFF))
			{
				OnNotify("Received data block # is not consecutive --> aborting transfer: ", fileTransfer, log4net.Core.Level.Error);
				SendError(fileTransfer, ErrorCode.ERROR_UNKNOWN_TID);
				fileTransfer.Close();
				fileTransfers.Remove(fileTransfer);
				return;
			}

			/// Update the current block number, save and acknowledge data
			fileTransfer.CurrentBlock++;
			for (int i = TftpHeaderLength; i < data.Length; i++) fileTransfer.StreamWriter.Write(data[i]);
			SendAck(fileTransfer);

			/// Close the file transfer if the block was the last one (data size < 512)
			if (data.Length < TftpDataBlockLength + TftpHeaderLength)
			{
				OnNotify("Download completed: ", fileTransfer, log4net.Core.Level.Info);
				fileTransfer.Close();
				fileTransfers.Remove(fileTransfer);
			}
		}

		#endregion


		/// <summary>
		/// Send part of a stream
		/// </summary>
		/// <param name="endpoint">location to send stream to</param>
		/// <param name="BlockNumber">512 byte block to send</param>
		private void SendError(FileTransfer fileTransfer, ErrorCode errorCode)
		{
			OnNotify("Sending Error Packet: errorCode = " + errorCode.ToString() + ", ", fileTransfer, log4net.Core.Level.Error);

			const int ErrorPacketLength = 10;

			Byte[] buffer = new Byte[ErrorPacketLength];

			buffer[0] = 0;
			buffer[1] = (byte)Opcode.ACK;
			buffer[2] = (byte)((fileTransfer.CurrentBlock & 0xFF00) / 256);
			buffer[3] = (byte)(fileTransfer.CurrentBlock & 0x00FF);
			buffer[4] = (Byte)'E';
			buffer[5] = (Byte)'r';
			buffer[6] = (Byte)'r';
			buffer[7] = (Byte)'o';
			buffer[8] = (Byte)'r';
			buffer[9] = 0;

			int ecode = udpClient.Send(buffer, ErrorPacketLength, fileTransfer.Endpoint);
			log.InfoFormat("Sent packet {0} ecode={1}", GetDebugMessage(buffer, ErrorPacketLength), ecode);

			if (ecode <= 0)
			{
				Debug.WriteLine("Error in send : {0}", ecode);
			}
		}

		/// <summary>
		/// Parse an error response
		/// </summary>
		/// <param name="data">data from packet</param>
		/// <param name="endpoint">client</param>
		private void ProcessError(Byte[] data, FileTransfer fileTransfer)
		{
			int errorCode = 256 * (int)data[2] + (int)data[3];
			string[] strData = Encoding.ASCII.GetString(data, 2, data.Length - 2).Split(Delimiter, 3);
			string message = strData[0];
			OnNotify("Received Error Packet: errorCode=" + errorCode.ToString() + ", message=" + message + ", ",
					 fileTransfer, log4net.Core.Level.Info);
		}

		#endregion
	}
}