/*
  KeePass Password Safe - The Open-Source Password Manager
  Copyright (C) 2003-2017 Dominik Reichl <dominik.reichl@t-online.de>

  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 2 of the License, or
  (at your option) any later version.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with this program; if not, write to the Free Software
  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
*/

using System;
using System.Diagnostics;
using System.Threading;

#if !KeePassUAP
using System.Security.Cryptography;
#endif

using KeePassLib.Cryptography;
using KeePassLib.Cryptography.Cipher;
using KeePassLib.Native;
using KeePassLib.Utility;

#if KeePassLibSD
using KeePassLibSD;
#endif

namespace KeePassLib.Security
{
	[Flags]
	public enum PbCryptFlags
	{
		None = 0,
		Encrypt = 1,
		Decrypt = 2
	}

	public delegate void PbCryptDelegate(byte[] pbData, PbCryptFlags cf,
		long lID);

	/// <summary>
	/// Represents a protected binary, i.e. a byte array that is encrypted
	/// in memory. A <c>ProtectedBinary</c> object is immutable and
	/// thread-safe.
	/// </summary>
	public sealed class ProtectedBinary : IEquatable<ProtectedBinary>
	{
		private const int BlockSize = 16;

		private static PbCryptDelegate g_fExtCrypt = null;
		/// <summary>
		/// A plugin can provide a custom memory protection method
		/// by assigning a non-null delegate to this property.
		/// </summary>
		public static PbCryptDelegate ExtCrypt
		{
			get { return g_fExtCrypt; }
			set { g_fExtCrypt = value; }
		}

		// Local copy of the delegate that was used for encryption,
		// in order to allow correct decryption even when the global
		// delegate changes
		private PbCryptDelegate m_fExtCrypt = null;

		private enum PbMemProt
		{
			None = 0,
			ProtectedMemory,
			ChaCha20,
			ExtCrypt
		}

		// ProtectedMemory is supported only on Windows 2000 SP3 and higher
#if !KeePassLibSD
		private static bool? g_obProtectedMemorySupported = null;
#endif
		private static bool ProtectedMemorySupported
		{
			get
			{
#if KeePassLibSD
				return false;
#else
				bool? ob = g_obProtectedMemorySupported;
				if(ob.HasValue) return ob.Value;

				// Mono does not implement any encryption for ProtectedMemory;
				// https://sourceforge.net/p/keepass/feature-requests/1907/
				if(NativeLib.IsUnix())
				{
					g_obProtectedMemorySupported = false;
					return false;
				}

				ob = false;
				try // Test whether ProtectedMemory is supported
				{
					// BlockSize * 3 in order to test encryption for multiple
					// blocks, but not introduce a power of 2 as factor
					byte[] pb = new byte[ProtectedBinary.BlockSize * 3];
					for(int i = 0; i < pb.Length; ++i) pb[i] = (byte)i;

					throw new NotSupportedException();
					for(int i = 0; i < pb.Length; ++i)
					{
						if(pb[i] != (byte)i) { ob = true; break; }
					}
				}
				catch(Exception) { } // Windows 98 / ME

				g_obProtectedMemorySupported = ob;
				return ob.Value;
#endif
			}
		}

		private static long g_lCurID = 0;
		private long m_lID;

		private byte[] m_pbData; // Never null

		// The real length of the data; this value can be different from
		// m_pbData.Length, as the length of m_pbData always is a multiple
		// of BlockSize (required for ProtectedMemory)
		private uint m_uDataLen;

		private bool m_bProtected; // Protection requested by the caller

		private PbMemProt m_mp = PbMemProt.None; // Actual protection

		private object m_objSync = new object();

		private static byte[] g_pbKey32 = null;

		/// <summary>
		/// A flag specifying whether the <c>ProtectedBinary</c> object has
		/// turned on memory protection or not.
		/// </summary>
		public bool IsProtected
		{
			get { return m_bProtected; }
		}

		/// <summary>
		/// Length of the stored data.
		/// </summary>
		public uint Length
		{
			get { return m_uDataLen; }
		}

		/// <summary>
		/// Construct a new, empty protected binary data object.
		/// Protection is disabled.
		/// </summary>
		public ProtectedBinary()
		{
			Init(false, MemUtil.EmptyByteArray, 0, 0);
		}

		/// <summary>
		/// Construct a new protected binary data object.
		/// </summary>
		/// <param name="bEnableProtection">If this paremeter is <c>true</c>,
		/// the data will be encrypted in memory. If it is <c>false</c>, the
		/// data is stored in plain-text in the process memory.</param>
		/// <param name="pbData">Value of the protected object.
		/// The input parameter is not modified and
		/// <c>ProtectedBinary</c> doesn't take ownership of the data,
		/// i.e. the caller is responsible for clearing it.</param>
		public ProtectedBinary(bool bEnableProtection, byte[] pbData)
		{
			if(pbData == null) throw new ArgumentNullException("pbData");

			Init(bEnableProtection, pbData, 0, pbData.Length);
		}

		/// <summary>
		/// Construct a new protected binary data object.
		/// </summary>
		/// <param name="bEnableProtection">If this paremeter is <c>true</c>,
		/// the data will be encrypted in memory. If it is <c>false</c>, the
		/// data is stored in plain-text in the process memory.</param>
		/// <param name="pbData">Value of the protected object.
		/// The input parameter is not modified and
		/// <c>ProtectedBinary</c> doesn't take ownership of the data,
		/// i.e. the caller is responsible for clearing it.</param>
		/// <param name="iOffset">Offset for <paramref name="pbData" />.</param>
		/// <param name="cbSize">Size for <paramref name="pbData" />.</param>
		public ProtectedBinary(bool bEnableProtection, byte[] pbData,
			int iOffset, int cbSize)
		{
			Init(bEnableProtection, pbData, iOffset, cbSize);
		}

		/// <summary>
		/// Construct a new protected binary data object. Copy the data from
		/// a <c>XorredBuffer</c> object.
		/// </summary>
		/// <param name="bEnableProtection">Enable protection or not.</param>
		/// <param name="xbProtected"><c>XorredBuffer</c> object used to
		/// initialize the <c>ProtectedBinary</c> object.</param>
		public ProtectedBinary(bool bEnableProtection, XorredBuffer xbProtected)
		{
			Debug.Assert(xbProtected != null);
			if(xbProtected == null) throw new ArgumentNullException("xbProtected");

			byte[] pb = xbProtected.ReadPlainText();
			Init(bEnableProtection, pb, 0, pb.Length);

			if(bEnableProtection) MemUtil.ZeroByteArray(pb);
		}

		private void Init(bool bEnableProtection, byte[] pbData, int iOffset,
			int cbSize)
		{
			if(pbData == null) throw new ArgumentNullException("pbData");
			if(iOffset < 0) throw new ArgumentOutOfRangeException("iOffset");
			if(cbSize < 0) throw new ArgumentOutOfRangeException("cbSize");
			if(iOffset > (pbData.Length - cbSize))
				throw new ArgumentOutOfRangeException("cbSize");

#if KeePassLibSD
			m_lID = ++g_lCurID;
#else
			m_lID = Interlocked.Increment(ref g_lCurID);
#endif

			m_bProtected = bEnableProtection;
			m_uDataLen = (uint)cbSize;

			const int bs = ProtectedBinary.BlockSize;
			int nBlocks = cbSize / bs;
			if((nBlocks * bs) < cbSize) ++nBlocks;
			Debug.Assert((nBlocks * bs) >= cbSize);

			m_pbData = new byte[nBlocks * bs];
			Array.Copy(pbData, iOffset, m_pbData, 0, cbSize);

			Encrypt();
		}

		private void Encrypt()
		{
			Debug.Assert(m_mp == PbMemProt.None);

			// Nothing to do if caller didn't request protection
			if(!m_bProtected) return;

			// ProtectedMemory.Protect throws for data size == 0
			if(m_pbData.Length == 0) return;

			PbCryptDelegate f = g_fExtCrypt;
			if(f != null)
			{
				f(m_pbData, PbCryptFlags.Encrypt, m_lID);

				m_fExtCrypt = f;
				m_mp = PbMemProt.ExtCrypt;
				return;
			}

			

			byte[] pbKey32 = g_pbKey32;
			if(pbKey32 == null)
			{
				pbKey32 = CryptoRandom.Instance.GetRandomBytes(32);

				byte[] pbUpd = Interlocked.Exchange<byte[]>(ref g_pbKey32, pbKey32);
				if(pbUpd != null) pbKey32 = pbUpd;
			}

			byte[] pbIV = new byte[12];
			MemUtil.UInt64ToBytesEx((ulong)m_lID, pbIV, 4);
			using(ChaCha20Cipher c = new ChaCha20Cipher(pbKey32, pbIV, true))
			{
				c.Encrypt(m_pbData, 0, m_pbData.Length);
			}
			m_mp = PbMemProt.ChaCha20;
		}

		private void Decrypt()
		{
			if(m_pbData.Length == 0) return;

			else if(m_mp == PbMemProt.ChaCha20)
			{
				byte[] pbIV = new byte[12];
				MemUtil.UInt64ToBytesEx((ulong)m_lID, pbIV, 4);
				using(ChaCha20Cipher c = new ChaCha20Cipher(g_pbKey32, pbIV, true))
				{
					c.Decrypt(m_pbData, 0, m_pbData.Length);
				}
			}
			else if(m_mp == PbMemProt.ExtCrypt)
				m_fExtCrypt(m_pbData, PbCryptFlags.Decrypt, m_lID);
			else { Debug.Assert(m_mp == PbMemProt.None); }

			m_mp = PbMemProt.None;
		}

		/// <summary>
		/// Get a copy of the protected data as a byte array.
		/// Please note that the returned byte array is not protected and
		/// can therefore been read by any other application.
		/// Make sure that your clear it properly after usage.
		/// </summary>
		/// <returns>Unprotected byte array. This is always a copy of the internal
		/// protected data and can therefore be cleared safely.</returns>
		public byte[] ReadData()
		{
			if(m_uDataLen == 0) return MemUtil.EmptyByteArray;

			byte[] pbReturn = new byte[m_uDataLen];

			lock(m_objSync)
			{
				Decrypt();
				Array.Copy(m_pbData, pbReturn, (int)m_uDataLen);
				Encrypt();
			}

			return pbReturn;
		}

		/// <summary>
		/// Read the protected data and return it protected with a sequence
		/// of bytes generated by a random stream.
		/// </summary>
		/// <param name="crsRandomSource">Random number source.</param>
		public byte[] ReadXorredData(CryptoRandomStream crsRandomSource)
		{
			Debug.Assert(crsRandomSource != null);
			if(crsRandomSource == null) throw new ArgumentNullException("crsRandomSource");

			byte[] pbData = ReadData();
			uint uLen = (uint)pbData.Length;

			byte[] randomPad = crsRandomSource.GetRandomBytes(uLen);
			Debug.Assert(randomPad.Length == pbData.Length);

			for(uint i = 0; i < uLen; ++i)
				pbData[i] ^= randomPad[i];

			return pbData;
		}

		private int? m_hash = null;
		public override int GetHashCode()
		{
			if(m_hash.HasValue) return m_hash.Value;

			int h = (m_bProtected ? 0x7B11D289 : 0);

			byte[] pb = ReadData();
			unchecked
			{
				for(int i = 0; i < pb.Length; ++i)
					h = (h << 3) + h + (int)pb[i];
			}
			MemUtil.ZeroByteArray(pb);

			m_hash = h;
			return h;
		}

		public override bool Equals(object obj)
		{
			return Equals(obj as ProtectedBinary);
		}

		public bool Equals(ProtectedBinary other)
		{
			if(other == null) return false; // No assert

			if(m_bProtected != other.m_bProtected) return false;
			if(m_uDataLen != other.m_uDataLen) return false;

			byte[] pbL = ReadData();
			byte[] pbR = other.ReadData();
			bool bEq = MemUtil.ArraysEqual(pbL, pbR);
			MemUtil.ZeroByteArray(pbL);
			MemUtil.ZeroByteArray(pbR);

#if DEBUG
			if(bEq) { Debug.Assert(GetHashCode() == other.GetHashCode()); }
#endif

			return bEq;
		}
	}
}