//	Spludlow Software
//	Copyright © Samuel P. Ludlow 2020 All Rights Reserved
//	Distributed under the terms of the GNU General Public License version 3
//	Distributed WITHOUT ANY WARRANTY; without implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE
//	https://www.spludlow.co.uk/LICENCE.TXT
//	The Spludlow logo is a registered trademark of Samuel P. Ludlow and may not be used without permission
//	v1.14

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

using System.Runtime.InteropServices;
using System.Collections;
using System.ComponentModel;

namespace Spludlow.Net
{
	public class DnsMx
	{
		[DllImport("dnsapi", EntryPoint = "DnsQuery_W", CharSet = CharSet.Unicode, SetLastError = true, ExactSpelling = true)]
		private static extern int DnsQuery([MarshalAs(UnmanagedType.VBByRefStr)]ref string pszName, QueryTypes wType, QueryOptions options, int aipServers, ref IntPtr ppQueryResults, int pReserved);

		[DllImport("dnsapi", CharSet = CharSet.Auto, SetLastError = true)]
		private static extern void DnsRecordListFree(IntPtr pRecordList, int FreeType);

		private enum QueryOptions
		{
			DNS_QUERY_ACCEPT_TRUNCATED_RESPONSE = 1,
			DNS_QUERY_BYPASS_CACHE = 8,
			DNS_QUERY_DONT_RESET_TTL_VALUES = 0x100000,
			DNS_QUERY_NO_HOSTS_FILE = 0x40,
			DNS_QUERY_NO_LOCAL_NAME = 0x20,
			DNS_QUERY_NO_NETBT = 0x80,
			DNS_QUERY_NO_RECURSION = 4,
			DNS_QUERY_NO_WIRE_QUERY = 0x10,
			DNS_QUERY_RESERVED = -16777216,
			DNS_QUERY_RETURN_MESSAGE = 0x200,
			DNS_QUERY_STANDARD = 0,
			DNS_QUERY_TREAT_AS_FQDN = 0x1000,
			DNS_QUERY_USE_TCP_ONLY = 2,
			DNS_QUERY_WIRE_ONLY = 0x100
		}

		private enum QueryTypes
		{
			DNS_TYPE_MX = 15
		}

		[StructLayout(LayoutKind.Sequential)]
		private struct MXRecord
		{
			public IntPtr pNext;
			public string pName;
			public short wType;
			public short wDataLength;
			public int flags;
			public int dwTtl;
			public int dwReserved;
			public IntPtr pNameExchange;
			public short wPreference;
			public short Pad;
		}

		public static string[][] GetMXRecordList(string domain)
		{
			Dictionary<int, List<string>> list = GetMXRecords(domain);

			List<int> preferences = new List<int>(list.Keys);
			preferences.Sort();

			List<string[]> result = new List<string[]>();

			foreach (int preference in preferences)
			{
				list[preference].Sort();		//	Should be left random for normal queries, but we want the same each time for identifing

				result.Add(list[preference].ToArray());
			}

			return result.ToArray();
		}

		public static Dictionary<int, List<string>> GetMXRecords(string domain)
		{
			Dictionary<int, List<string>> list = new Dictionary<int, List<string>>();

			IntPtr startRecord = IntPtr.Zero;

			bool keepTrying = true;
			for (int retry = 0; retry < 3 && keepTrying == true; ++retry)
			{
				int result;
				if ((result = DnsMx.DnsQuery(ref domain, QueryTypes.DNS_TYPE_MX, QueryOptions.DNS_QUERY_BYPASS_CACHE, 0, ref startRecord, 0)) != 0)
				{
					if (result == 0x0000232B)   //	DNS name does not exist
						return list;

					if (result == 0x0000251D)   //	No records found for given DNS query
						return list;

					Win32Exception exception = new Win32Exception(result);

					keepTrying = false;

					//	This operation returned because the timeout period expired 000005B4
					//	DNS server failure 0000232A
					if (result == 0x000005B4 || result == 0x0000232A)
					{
						keepTrying = true;
						Spludlow.Log.Warning("GetMXRecords; Retry: " + retry + ", " + exception.Message + " " + result.ToString("X8") + ": " + domain, exception);
						System.Threading.Thread.Sleep(3 * 1000);
					}
					else
					{
						throw new ApplicationException("GetMXRecords, DnsQuery; " + exception.Message + " " + result.ToString("X8") + ": " + domain, exception);
					}
				}
			}

			if (startRecord == IntPtr.Zero)
				return list;

			try
			{
				MXRecord mxRecord;
				for (IntPtr currentRecord = startRecord; !currentRecord.Equals(IntPtr.Zero); currentRecord = mxRecord.pNext)
				{
					mxRecord = (MXRecord)Marshal.PtrToStructure(currentRecord, typeof(MXRecord));
					if (mxRecord.wType == 15)   //	Can also have IPv6 address with same name and preference
					{
						int preference = mxRecord.wPreference;
						string name = Marshal.PtrToStringAuto(mxRecord.pNameExchange);

						if (list.ContainsKey(preference) == false)
							list.Add(preference, new List<string>());

						list[preference].Add(name);
					}
				}

				return list;
			}
			finally
			{
				DnsMx.DnsRecordListFree(startRecord, 0);
			}
		}
	}
}