using System;
using System.Text;
using System.Runtime.InteropServices;

namespace AALabs.Utilities
{
	/// <summary>
	/// An interface to the protected storage service in Windows.  You can use this
	/// to store passwords, keys, codes, etc, in at least a slightly more secure 
	/// fashion than the registry or disk.  
	/// <para><a href="http://msdn.microsoft.com/architecture/application/default.aspx?pull=/library/en-us/dnnetsec/html/SecNetHT07.asp">
	/// http://msdn.microsoft.com/architecture/application/default.aspx?pull=/library/en-us/dnnetsec/html/SecNetHT07.asp</a></para>
	/// </summary>
	public class DataProtection
	{
		#region DLL Imports
		[DllImport("Crypt32.dll", SetLastError=true,
			 CharSet=System.Runtime.InteropServices.CharSet.Auto)]
		private static extern bool CryptProtectData(
			ref DATA_BLOB pDataIn, 
			string szDataDescr, 
			ref DATA_BLOB pOptionalEntropy,
			IntPtr pvReserved, 
			ref CRYPTPROTECT_PROMPTSTRUCT pPromptStruct, 
			int dwFlags,
			ref DATA_BLOB pDataOut);
		[DllImport("Crypt32.dll", SetLastError=true, 
			 CharSet=System.Runtime.InteropServices.CharSet.Auto)]
		private static extern bool CryptUnprotectData(
			ref DATA_BLOB pDataIn, 
			ref string szDataDescr, 
			ref DATA_BLOB pOptionalEntropy, 
			IntPtr pvReserved, 
			ref CRYPTPROTECT_PROMPTSTRUCT 
			pPromptStruct, 
			int dwFlags, 
			ref DATA_BLOB pDataOut);
		/*
		[DllImport("kernel32.dll", 
			 CharSet=System.Runtime.InteropServices.CharSet.Auto)]
		private unsafe static extern int FormatMessage(int dwFlags, 
			ref IntPtr lpSource, 
			int dwMessageId,
			int dwLanguageId, 
			ref String lpBuffer, 
			int nSize,
			IntPtr *Arguments);
			*/
		#endregion

		#region DLL Structs
		[StructLayout(LayoutKind.Sequential, CharSet=CharSet.Unicode)]
			internal struct DATA_BLOB
		{
			public int cbData;
			public IntPtr pbData;
		}

		[StructLayout(LayoutKind.Sequential, CharSet=CharSet.Unicode)]
			internal struct CRYPTPROTECT_PROMPTSTRUCT
		{
			public int cbSize;
			public int dwPromptFlags;
			public IntPtr hwndApp;
			public String szPrompt;
		}
		static private IntPtr NullPtr = ((IntPtr)((int)(0)));
		private const int CRYPTPROTECT_UI_FORBIDDEN = 0x1;
		private const int CRYPTPROTECT_LOCAL_MACHINE = 0x4;
		#endregion

		/// <summary>
		/// Which data store to use
		/// </summary>
		public enum Store 
		{
			/// <summary>
			/// Machine store does not include user info, but can't be used on another machine easily.
			/// </summary>
			USE_MACHINE_STORE = 1, 
			/// <summary>
			/// Both machine and user specific data are included.
			/// </summary>
			USE_USER_STORE
		};

		private Store store;

		/// <summary>
		/// Construct a data protector that stores data in the specified protected store.
		/// </summary>
		/// <param name="tStore">Machine specific store or user+machine specific store.</param>
		public DataProtection(Store tStore)
		{
			store = tStore;
		}

		/// <summary>
		/// Encrypt a plaintext value given a salt.  A machine specific (and
		/// possibly user specific if the user store is selected) key is used
		/// for the encryption.
		/// </summary>
		/// <param name="plainText">ASCII string</param>
		/// <param name="entropy">ASCII entropy</param>
		/// <returns>Base64 encoded string</returns>
		public string Encrypt(string plainText, string entropy)
		{
			byte[] arr = Encrypt(Encoding.ASCII.GetBytes(plainText),
				entropy != null ? Encoding.ASCII.GetBytes(entropy) : null);
			return Convert.ToBase64String(arr);
		}

		/// <summary>
		/// Encrypt a plaintext value given a salt.  A machine specific (and
		/// possibly user specific if the user store is selected) key is used
		/// for the encryption.
		/// </summary>
		/// <param name="plainText">Original plain text</param>
		/// <param name="optionalEntropy">Optional salt</param>
		/// <returns>byte array of ciphertext</returns>
		public byte[] Encrypt(byte[] plainText, byte[] optionalEntropy)
		{
			bool retVal = false;

			DATA_BLOB plainTextBlob = new DATA_BLOB();
			DATA_BLOB cipherTextBlob = new DATA_BLOB();
			DATA_BLOB entropyBlob = new DATA_BLOB();

			CRYPTPROTECT_PROMPTSTRUCT prompt = new 
				CRYPTPROTECT_PROMPTSTRUCT();
			InitPromptstruct(ref prompt);

			int dwFlags;
			try
			{
				try
				{
					int bytesSize = plainText.Length;
					plainTextBlob.pbData = Marshal.AllocHGlobal(bytesSize);
					if(IntPtr.Zero == plainTextBlob.pbData)
					{
						throw new ApplicationException("DP.PlaintextError");
					}
					plainTextBlob.cbData = bytesSize;
					Marshal.Copy(plainText, 0, plainTextBlob.pbData, bytesSize);
				}
				catch(Exception ex)
				{
					throw new ApplicationException("DP.MarshalError",ex);
				}
				if(Store.USE_MACHINE_STORE == store)
				{//Using the machine store, should be providing entropy.
					dwFlags = 
						CRYPTPROTECT_LOCAL_MACHINE|CRYPTPROTECT_UI_FORBIDDEN;
					//Check to see if the entropy is null
					if(null == optionalEntropy)
					{//Allocate something
						optionalEntropy = new byte[0];
					}
					try
					{
						int bytesSize = optionalEntropy.Length;
						entropyBlob.pbData = 
							Marshal.AllocHGlobal(optionalEntropy.Length);;
						if(IntPtr.Zero == entropyBlob.pbData)
						{
							throw new Exception("Unable to allocate entropy data buffer.");
						}
						Marshal.Copy(optionalEntropy, 0, entropyBlob.pbData, 
							bytesSize);
						entropyBlob.cbData = bytesSize;
					}
					catch(Exception ex)
					{
						throw new Exception("Exception entropy marshalling data. " + 
							ex.Message);
					}
				}
				else
				{//Using the user store
					dwFlags = CRYPTPROTECT_UI_FORBIDDEN;
				}
				retVal = CryptProtectData(ref plainTextBlob, "WAFData", ref 
					entropyBlob, 
					IntPtr.Zero, ref prompt, dwFlags, 
					ref cipherTextBlob);
				if(false == retVal)
				{
					throw new Exception("Encryption failed. " + 
						GetErrorMessage(Marshal.GetLastWin32Error()));
				}
				//Free the blob and entropy.
				/*
				if(IntPtr.Zero != cipherTextBlob.pbData)
				{
					Marshal.FreeHGlobal(cipherTextBlob.pbData);
				}
				if(IntPtr.Zero != entropyBlob.pbData)
				{
					Marshal.FreeHGlobal(entropyBlob.pbData);
				}
				*/
			}
			catch(Exception ex)
			{
				throw new Exception("Exception encrypting. " + ex.Message);
			}
			byte[] cipherText = new byte[cipherTextBlob.cbData];
			Marshal.Copy(cipherTextBlob.pbData, cipherText, 0, 
				cipherTextBlob.cbData);
			return cipherText;
		}

		/// <summary>
		/// Decrypt a ciphertext value given a salt.  A machine specific (and
		/// possibly user specific if the user store is selected) key was used
		/// for the encryption.
		/// </summary>
		/// <param name="cipherText">Base64 encoding string</param>
		/// <param name="optionalEntropy">Optional salt</param>
		public string Decrypt(string cipherText, string optionalEntropy)
		{
			byte[] arr = Decrypt(Convert.FromBase64String(cipherText),
				optionalEntropy != null ? Encoding.ASCII.GetBytes(optionalEntropy) : null);
			return Encoding.ASCII.GetString(arr);
		}

		/// <summary>
		/// Decrypt a ciphertext value given a salt.  A machine specific (and
		/// possibly user specific if the user store is selected) key was used
		/// for the encryption.
		/// </summary>
		/// <param name="cipherText">Byte array encrypted with Encrypt</param>
		/// <param name="optionalEntropy">Optional salt</param>
		public byte[] Decrypt(byte[] cipherText, byte[] optionalEntropy)
		{
			bool retVal = false;
			DATA_BLOB plainTextBlob = new DATA_BLOB();
			DATA_BLOB cipherBlob = new DATA_BLOB();
			CRYPTPROTECT_PROMPTSTRUCT prompt = new 
				CRYPTPROTECT_PROMPTSTRUCT();
			InitPromptstruct(ref prompt);
			try
			{
				try
				{
					int cipherTextSize = cipherText.Length;
					cipherBlob.pbData = Marshal.AllocHGlobal(cipherTextSize);
					if(IntPtr.Zero == cipherBlob.pbData)
					{
						throw new Exception("Unable to allocate cipherText buffer.");
					}
					cipherBlob.cbData = cipherTextSize;
					Marshal.Copy(cipherText, 0, cipherBlob.pbData, 
						cipherBlob.cbData);
				}
				catch(Exception ex)
				{
					throw new Exception("Exception marshalling data. " + 
						ex.Message);
				}
				DATA_BLOB entropyBlob = new DATA_BLOB();
				int dwFlags;
				if(Store.USE_MACHINE_STORE == store)
				{//Using the machine store, should be providing entropy.
					dwFlags = 
						CRYPTPROTECT_LOCAL_MACHINE|CRYPTPROTECT_UI_FORBIDDEN;
					//Check to see if the entropy is null
					if(null == optionalEntropy)
					{//Allocate something
						optionalEntropy = new byte[0];
					}
					try
					{
						int bytesSize = optionalEntropy.Length;
						entropyBlob.pbData = Marshal.AllocHGlobal(bytesSize);
						if(IntPtr.Zero == entropyBlob.pbData)
						{
							throw new Exception("Unable to allocate entropy buffer.");
						}
						entropyBlob.cbData = bytesSize;
						Marshal.Copy(optionalEntropy, 0, entropyBlob.pbData, 
							bytesSize);
					}
					catch(Exception ex)
					{
						throw new Exception("Exception entropy marshalling data. " + 
							ex.Message);
					}
				}
				else
				{//Using the user store
					dwFlags = CRYPTPROTECT_UI_FORBIDDEN;
				}
				string desc = String.Empty;
				retVal = CryptUnprotectData(ref cipherBlob, ref desc, ref 
					entropyBlob, 
					IntPtr.Zero, ref prompt, dwFlags, 
					ref plainTextBlob);
				if(false == retVal)
				{
					throw new Exception("Decryption failed. " + 
						GetErrorMessage(Marshal.GetLastWin32Error()));
				}
				//Free the blob and entropy.
				if(IntPtr.Zero != cipherBlob.pbData)
				{
					Marshal.FreeHGlobal(cipherBlob.pbData);
				}
				if(IntPtr.Zero != entropyBlob.pbData)
				{
					Marshal.FreeHGlobal(entropyBlob.pbData);
				}
			}
			catch(Exception ex)
			{
				throw new Exception("Exception decrypting. " + ex.Message);
			}
			byte[] plainText = new byte[plainTextBlob.cbData];
			Marshal.Copy(plainTextBlob.pbData, plainText, 0, 
				plainTextBlob.cbData);
			return plainText;
		}

		private void InitPromptstruct(ref CRYPTPROTECT_PROMPTSTRUCT ps) 
		{
			ps.cbSize = Marshal.SizeOf(typeof(CRYPTPROTECT_PROMPTSTRUCT));
			ps.dwPromptFlags = 0;
			ps.hwndApp = NullPtr;
			ps.szPrompt = null;
		}

		private static string GetErrorMessage(int errorCode)
		{
			return String.Format("DP.ErrorCode", errorCode);
		}

	}
}