using System;
using System.IO;
using System.Security.Cryptography;

namespace Synercoding.FormsAuthentication.Encryption
{
    // Based upon: https://github.com/Microsoft/referencesource/blob/master/System.Web/Security/Cryptography/NetFXCryptoService.cs
    internal sealed class NetFXCryptoService : ICryptoService
    {
        private readonly ICryptoAlgorithmFactory _cryptoAlgorithmFactory;
        private readonly CryptographicKey _encryptionKey;
        private readonly bool _predictableIV;
        private readonly CryptographicKey _validationKey;

        public NetFXCryptoService(ICryptoAlgorithmFactory cryptoAlgorithmFactory, CryptographicKey encryptionKey, CryptographicKey validationKey, bool predictableIV = false)
        {
            _cryptoAlgorithmFactory = cryptoAlgorithmFactory;
            _encryptionKey = encryptionKey;
            _validationKey = validationKey;
            _predictableIV = predictableIV;
        }

        public byte[] Protect(byte[] clearData)
        {
            // The entire operation is wrapped in a 'checked' block because any overflows should be treated as failures.
            checked
            {

                // These SymmetricAlgorithm instances are single-use; we wrap it in a 'using' block.
                using (SymmetricAlgorithm encryptionAlgorithm = _cryptoAlgorithmFactory.GetEncryptionAlgorithm())
                {
                    // Initialize the algorithm with the specified key and an appropriate IV
                    encryptionAlgorithm.Key = _encryptionKey.GetKeyMaterial();

                    if (_predictableIV)
                    {
                        // The caller wanted the output to be predictable (e.g. for caching), so we'll create an
                        // appropriate IV directly from the input buffer. The IV length is equal to the block size.
                        encryptionAlgorithm.IV = CryptoUtil.CreatePredictableIV(clearData, encryptionAlgorithm.BlockSize);
                    }
                    else
                    {
                        // If the caller didn't ask for a predictable IV, just let the algorithm itself choose one.
                        encryptionAlgorithm.GenerateIV();
                    }
                    byte[] iv = encryptionAlgorithm.IV;

                    using (MemoryStream memStream = new MemoryStream())
                    {
                        memStream.Write(iv, 0, iv.Length);

                        // At this point:
                        // memStream := IV

                        // Write the encrypted payload to the memory stream.
                        using (ICryptoTransform encryptor = encryptionAlgorithm.CreateEncryptor())
                        {
                            using (CryptoStream cryptoStream = new CryptoStream(memStream, encryptor, CryptoStreamMode.Write))
                            {
                                cryptoStream.Write(clearData, 0, clearData.Length);
                                cryptoStream.FlushFinalBlock();

                                // At this point:
                                // memStream := IV || Enc(Kenc, IV, clearData)

                                // These KeyedHashAlgorithm instances are single-use; we wrap it in a 'using' block.
                                using (KeyedHashAlgorithm signingAlgorithm = _cryptoAlgorithmFactory.GetValidationAlgorithm())
                                {
                                    // Initialize the algorithm with the specified key
                                    signingAlgorithm.Key = _validationKey.GetKeyMaterial();

                                    // Compute the signature
                                    ArraySegment<byte> segment;
                                    memStream.TryGetBuffer(out segment);
                                    byte[] signature = signingAlgorithm.ComputeHash(segment.Array, 0, (int)memStream.Length);

                                    // At this point:
                                    // memStream := IV || Enc(Kenc, IV, clearData)
                                    // signature := Sign(Kval, IV || Enc(Kenc, IV, clearData))

                                    // Append the signature to the encrypted payload
                                    memStream.Write(signature, 0, signature.Length);

                                    // At this point:
                                    // memStream := IV || Enc(Kenc, IV, clearData) || Sign(Kval, IV || Enc(Kenc, IV, clearData))

                                    // Algorithm complete
                                    byte[] protectedData = memStream.ToArray();
                                    return protectedData;
                                }
                            }
                        }
                    }
                }
            }
        }

        public byte[] Unprotect(byte[] protectedData)
        {
            // The entire operation is wrapped in a 'checked' block because any overflows should be treated as failures.
            checked
            {

                // We want to check that the input is in the form:
                // protectedData := IV || Enc(Kenc, IV, clearData) || Sign(Kval, IV || Enc(Kenc, IV, clearData))

                // Definitions used in this method:
                // encryptedPayload := Enc(Kenc, IV, clearData)
                // signature := Sign(Kval, IV || encryptedPayload)

                // These SymmetricAlgorithm instances are single-use; we wrap it in a 'using' block.
                using (SymmetricAlgorithm decryptionAlgorithm = _cryptoAlgorithmFactory.GetEncryptionAlgorithm())
                {
                    decryptionAlgorithm.Key = _encryptionKey.GetKeyMaterial();

                    // These KeyedHashAlgorithm instances are single-use; we wrap it in a 'using' block.
                    using (KeyedHashAlgorithm validationAlgorithm = _cryptoAlgorithmFactory.GetValidationAlgorithm())
                    {
                        validationAlgorithm.Key = _validationKey.GetKeyMaterial();

                        // First, we need to verify that protectedData is even long enough to contain
                        // the required components (IV, encryptedPayload, signature).

                        int ivByteCount = decryptionAlgorithm.BlockSize / 8; // IV length is equal to the block size
                        int signatureByteCount = validationAlgorithm.HashSize / 8;
                        int encryptedPayloadByteCount = protectedData.Length - ivByteCount - signatureByteCount;
                        if (encryptedPayloadByteCount <= 0)
                        {
                            // protectedData doesn't meet minimum length requirements
                            return null;
                        }

                        // If that check passes, we need to detect payload tampering.

                        // Compute the signature over the IV and encrypted payload
                        // computedSignature := Sign(Kval, IV || encryptedPayload)
                        byte[] computedSignature = validationAlgorithm.ComputeHash(protectedData, 0, ivByteCount + encryptedPayloadByteCount);

                        if (!CryptoUtil.BuffersAreEqual(
                            buffer1: protectedData, buffer1Offset: ivByteCount + encryptedPayloadByteCount, buffer1Count: signatureByteCount,
                            buffer2: computedSignature, buffer2Offset: 0, buffer2Count: computedSignature.Length))
                        {

                            // the computed signature didn't match the incoming signature, which is a sign of payload tampering
                            return null;
                        }

                        // At this point, we're certain that we generated the signature over this payload,
                        // so we can go ahead with decryption.

                        // Populate the IV from the incoming stream
                        byte[] iv = new byte[ivByteCount];
                        Buffer.BlockCopy(protectedData, 0, iv, 0, iv.Length);
                        decryptionAlgorithm.IV = iv;

                        // Write the decrypted payload to the memory stream.
                        using (MemoryStream memStream = new MemoryStream())
                        {
                            using (ICryptoTransform decryptor = decryptionAlgorithm.CreateDecryptor())
                            {
                                using (CryptoStream cryptoStream = new CryptoStream(memStream, decryptor, CryptoStreamMode.Write))
                                {
                                    cryptoStream.Write(protectedData, ivByteCount, encryptedPayloadByteCount);
                                    cryptoStream.FlushFinalBlock();

                                    // At this point
                                    // memStream := clearData

                                    byte[] clearData = memStream.ToArray();
                                    return clearData;
                                }
                            }
                        }
                    }
                }
            }
        }

    }
}