// Copyright (c) .NET Foundation. All rights reserved.
// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.

using System;
using System.Diagnostics;
using System.Security.Cryptography;
using System.Text;

namespace Microsoft.AspNet.Cryptography.KeyDerivation.PBKDF2
{
    /// <summary>
    /// A PBKDF2 provider which utilizes the managed hash algorithm classes as PRFs.
    /// This isn't the preferred provider since the implementation is slow, but it is provided as a fallback.
    /// </summary>
    internal sealed class ManagedPbkdf2Provider : IPbkdf2Provider
    {
        public byte[] DeriveKey(string password, byte[] salt, KeyDerivationPrf prf, int iterationCount, int numBytesRequested)
        {
            Debug.Assert(password != null);
            Debug.Assert(salt != null);
            Debug.Assert(iterationCount > 0);
            Debug.Assert(numBytesRequested > 0);

            // PBKDF2 is defined in NIST SP800-132, Sec. 5.3.
            // http://csrc.nist.gov/publications/nistpubs/800-132/nist-sp800-132.pdf

            byte[] retVal = new byte[numBytesRequested];
            int numBytesWritten = 0;
            int numBytesRemaining = numBytesRequested;

            // For each block index, U_0 := Salt || block_index
            byte[] saltWithBlockIndex = new byte[checked(salt.Length + sizeof(uint))];
            Buffer.BlockCopy(salt, 0, saltWithBlockIndex, 0, salt.Length);

            using (var hashAlgorithm = PrfToManagedHmacAlgorithm(prf, password))
            {
                for (uint blockIndex = 1; numBytesRemaining > 0; blockIndex++)
                {
                    // write the block index out as big-endian
                    saltWithBlockIndex[saltWithBlockIndex.Length - 4] = (byte)(blockIndex >> 24);
                    saltWithBlockIndex[saltWithBlockIndex.Length - 3] = (byte)(blockIndex >> 16);
                    saltWithBlockIndex[saltWithBlockIndex.Length - 2] = (byte)(blockIndex >> 8);
                    saltWithBlockIndex[saltWithBlockIndex.Length - 1] = (byte)blockIndex;

                    // U_1 = PRF(U_0) = PRF(Salt || block_index)
                    // T_blockIndex = U_1
                    byte[] U_iter = hashAlgorithm.ComputeHash(saltWithBlockIndex); // this is U_1
                    byte[] T_blockIndex = U_iter;

                    for (int iter = 1; iter < iterationCount; iter++)
                    {
                        U_iter = hashAlgorithm.ComputeHash(U_iter);
                        XorBuffers(src: U_iter, dest: T_blockIndex);
                        // At this point, the 'U_iter' variable actually contains U_{iter+1} (due to indexing differences).
                    }

                    // At this point, we're done iterating on this block, so copy the transformed block into retVal.
                    int numBytesToCopy = Math.Min(numBytesRemaining, T_blockIndex.Length);
                    Buffer.BlockCopy(T_blockIndex, 0, retVal, numBytesWritten, numBytesToCopy);
                    numBytesWritten += numBytesToCopy;
                    numBytesRemaining -= numBytesToCopy;
                }
            }

            // retVal := T_1 || T_2 || ... || T_n, where T_n may be truncated to meet the desired output length
            return retVal;
        }

        private static KeyedHashAlgorithm PrfToManagedHmacAlgorithm(KeyDerivationPrf prf, string password)
        {
            byte[] passwordBytes = Encoding.UTF8.GetBytes(password);
            try
            {
                switch (prf)
                {
                    case KeyDerivationPrf.HMACSHA1:
                        return new HMACSHA1(passwordBytes);
                    case KeyDerivationPrf.HMACSHA256:
                        return new HMACSHA256(passwordBytes);
                    case KeyDerivationPrf.HMACSHA512:
                        return new HMACSHA512(passwordBytes);
                    default:
                        throw CryptoUtil.Fail("Unrecognized PRF.");
                }
            }
            finally
            {
                // The HMAC ctor makes a duplicate of this key; we clear original buffer to limit exposure to the GC.
                Array.Clear(passwordBytes, 0, passwordBytes.Length);
            }
        }

        private static void XorBuffers(byte[] src, byte[] dest)
        {
            // Note: dest buffer is mutated.
            Debug.Assert(src.Length == dest.Length);
            for (int i = 0; i < src.Length; i++)
            {
                dest[i] ^= src[i];
            }
        }
    }
}