using System;
using System.IO;
using System.Security.Cryptography;
using Org.BouncyCastle.Crypto.Engines;
using Org.BouncyCastle.Crypto.Modes;
using Org.BouncyCastle.Crypto.Parameters;

namespace Xylem.Common.Cryptology.Security
{
    /// <summary>
    /// Class for cryptology.
    /// </summary>
    public static class Cryptology
    {
        private const Int32 MacSize = 16;
        private const Int32 MacBitSize = MacSize * 8;

        /// <summary>
        /// AES-GCM Decryption using bouncy Castle nuget package
        /// NOTE - add nonce + salt as aad into GCM encryption/authentication
        /// </summary>
        /// <param name="saltSize"></param>
        /// <param name="nonceSize"></param>
        /// <param name="keySize"></param>
        /// <param name="primaryKey"> byte array of key components</param>
        /// <param name="encryptedData">  encrypted cipher byte array = nonce | salt | cipherText | MAC  </param>
        /// <returns> encrypted cipher byte array </returns>
        /// <remarks date="2021-Jan-21" author="Thomas Wiedebusch">
        /// - Initial based on code snipped from Ray Savarda.
        /// </remarks>
        public static Byte[] AesGcmDecryption(Int32 saltSize, Int32 nonceSize, Int32 keySize, Byte[] primaryKey, 
            Byte[] encryptedData)
        {
            // Extract nonce and salt from encryptedData
            var nonce = new Byte[nonceSize];
            Array.Copy(encryptedData, nonce, nonceSize);
            var salt = new Byte[saltSize];
            Array.Copy(encryptedData, nonceSize, salt, 0, saltSize);

            var aes256Key = BuildAes256Key(primaryKey, salt, keySize);
            
            //extract aad for GCM call
            var aad = new Byte[nonceSize + saltSize];
            Array.Copy(encryptedData, 0, aad, 0, aad.Length);

            var cipher = new GcmBlockCipher(new AesEngine());
            var parameters = new AeadParameters(new KeyParameter(aes256Key), MacBitSize, nonce, aad);
            cipher.Init(false, parameters);

            var gcmCipherText = new Byte[encryptedData.Length - aad.Length];
            Array.Copy(encryptedData, aad.Length, gcmCipherText, 0, encryptedData.Length - aad.Length);

            var rawData = new Byte[cipher.GetOutputSize(gcmCipherText.Length)];
            
            //Note aad handled by AadParameters call
            var len = cipher.ProcessBytes(gcmCipherText, 0, gcmCipherText.Length, rawData, 0); 
            
            //TODO THW Need to add exception handlers - if MAC failure, get exception!
            /*var gcmResultCount =*/ cipher.DoFinal(rawData, len);

            return rawData;
        }

        /// <summary>
        /// AES-GCM Encryption using bouncy Castle nuget package
        /// NOTE - add nonce + salt as aad into GCM encryption/authentication
        /// </summary>
        /// <param name="saltSize"></param>
        /// <param name="nonceSize"></param>
        /// <param name="keySize"></param>
        /// <param name="primaryKey"> byte array of key components</param>
        /// <param name="rawData"> raw data stream for encryption - payload </param>
        /// <returns> encrypted cipher byte array = nonce | salt | cipherText | MAC </returns>
        /// <remarks date="2021-Jan-21" author="Thomas Wiedebusch">
        /// - Initial based on code snipped from Ray Savarda.
        /// </remarks>
        public static Byte[] AesGcmEncryption(Int32 saltSize, Int32 nonceSize, Int32 keySize, Byte[] primaryKey, 
            Byte[] rawData)
        {
            // Initial salt generation
            var salt = BuildRandomValue(saltSize);
            var nonce = BuildRandomValue(nonceSize);
            var aes256Key = BuildAes256Key(primaryKey, salt, keySize);
            
            // put together the aad so it can be included in GCM call
            var aad = new Byte[nonceSize + saltSize];
            Array.Copy(nonce, 0, aad, 0, nonce.Length);
            Array.Copy(salt, 0, aad, nonce.Length, salt.Length);

            var cipher = new GcmBlockCipher(new AesEngine());
            //include AAD for MAC protection!
            var parameters = new AeadParameters(new KeyParameter(aes256Key), MacBitSize, nonce, aad);  
            cipher.Init(true, parameters);
            var cipherText = new Byte[cipher.GetOutputSize(rawData.Length)];
            var len = cipher.ProcessBytes(rawData, 0, rawData.Length, cipherText, 0);
            cipher.DoFinal(cipherText, len); //Note Cipher text includes MAC at end!

            //var noMac = new Byte[cipherText.Length - MacSize];
            //Array.Copy(cipherText, noMac, cipherText.Length - MacSize);

            //var onlyMac = new Byte[MacSize];
            //Array.ConstrainedCopy(cipherText, cipherText.Length - MacSize, onlyMac, 0, onlyMac.Length);

            //Assemble output memory stream = nonce + salt + cipherText + MAC 
            var outMs = new MemoryStream();
            using (var binaryWriter = new BinaryWriter(outMs))
            {
                //Prepend Nonce
                binaryWriter.Write(nonce);
                //add salt
                binaryWriter.Write(salt);
                //Write Cipher Text
                binaryWriter.Write(cipherText);
            }
           
            return outMs.ToArray();
        }

        /// <summary>
        /// Build random value
        /// </summary>
        /// <param name="size"></param>
        /// <returns> byte array of random value </returns>
        /// <remarks date="2021-Jan-21" author="Thomas Wiedebusch">
        /// - Initial based on code snipped from Ray Savarda.
        /// </remarks>
        public static Byte[] BuildRandomValue(Int32 size)
        {
            var random = new Byte[size];
            var rngCsp = new RNGCryptoServiceProvider();
            rngCsp.GetBytes(random);
            return random;
        }

        /// <summary>
        /// Build AES256 key
        /// </summary>
        /// <param name="primaryKey">primary information for key generation</param>
        /// <param name="salt"></param>
        /// <param name="aes256KeySize"></param>
        /// <returns> AES256 key </returns>
        /// <remarks date="2021-Jan-21" author="Thomas Wiedebusch">
        /// - Initial based on code snipped from Ray Savarda.
        /// </remarks>
        public static Byte[] BuildAes256Key(Byte[] primaryKey, Byte[] salt, Int32 aes256KeySize)
        {
            const Int32 iterations = 10000; //(typically 2,000 - 10,000)
            var aes256Key = new Rfc2898DeriveBytes(primaryKey, salt, iterations, HashAlgorithmName.SHA256);
            return aes256Key.GetBytes(aes256KeySize);
        }
    }
}