EonaCat.Transfer/EonaCat.Transfer/Models/AesEncryptionProvider.cs

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

namespace EonaCat.Transfer.Models
{
    public class AesEncryptionProvider : IDisposable
    {
        private readonly string _salt = "EonaCat_Transfer_Salt";
        private readonly Aes _aes;
        private readonly ICryptoTransform _encryptor;
        private readonly ICryptoTransform _decryptor;
        private bool _disposed;

        public byte[] Key => _aes.Key;
        public byte[] IV => _aes.IV;

        public AesEncryptionProvider(AesConfig config, string salt = "EonaCat_Transfer_Salt")
        {
            _salt = salt;
            _aes = Aes.Create();
            _aes.KeySize = (int)config.KeySize;
            _aes.Mode = config.Mode;
            _aes.Padding = config.Padding;

            if (config.Key != null && config.IV != null)
            {
                _aes.Key = config.Key;
                _aes.IV = config.IV;
            }
            else if (!string.IsNullOrEmpty(config.SharedSecret))
            {
                // Derive key from shared secret using PBKDF2
                DeriveKeyFromSecret(config.SharedSecret, config.KeySize);
            }
            else
            {
                // Generate random key and IV
                _aes.GenerateKey();
                _aes.GenerateIV();
            }

            _encryptor = _aes.CreateEncryptor();
            _decryptor = _aes.CreateDecryptor();
        }

        private void DeriveKeyFromSecret(string secret, AesKeySize keySize)
        {
            // Use PBKDF2 for key derivation
            var salt = Encoding.UTF8.GetBytes(_salt);

#if NET6_0_OR_GREATER
            using (var deriveBytes = new Rfc2898DeriveBytes(
                secret, 
                salt, 
                100000, 
                HashAlgorithmName.SHA256))
#else
            using (var deriveBytes = new Rfc2898DeriveBytes(secret, salt, 100000))
#endif
            {
                _aes.Key = deriveBytes.GetBytes((int)keySize / 8);
                _aes.IV = deriveBytes.GetBytes(16); // AES block size is always 128 bits (16 bytes)
            }
        }

        public byte[] Encrypt(byte[] data)
        {
            if (data == null || data.Length == 0)
            {
                return data;
            }

            using (var ms = new MemoryStream())
            using (var cs = new CryptoStream(ms, _encryptor, CryptoStreamMode.Write))
            {
                cs.Write(data, 0, data.Length);
                cs.FlushFinalBlock();
                return ms.ToArray();
            }
        }

        public byte[] Decrypt(byte[] encryptedData)
        {
            if (encryptedData == null || encryptedData.Length == 0)
            {
                return encryptedData;
            }

            using (var ms = new MemoryStream())
            using (var cs = new CryptoStream(ms, _decryptor, CryptoStreamMode.Write))
            {
                cs.Write(encryptedData, 0, encryptedData.Length);
                cs.FlushFinalBlock();
                return ms.ToArray();
            }
        }

        public async Task<byte[]> EncryptAsync(byte[] data)
        {
            return await Task.Run(() => Encrypt(data));
        }

        public async Task<byte[]> DecryptAsync(byte[] encryptedData)
        {
            return await Task.Run(() => Decrypt(encryptedData));
        }

        public static AesConfig GenerateConfig(AesKeySize keySize = AesKeySize.Aes256)
        {
            using (var aes = Aes.Create())
            {
                aes.KeySize = (int)keySize;
                aes.GenerateKey();
                aes.GenerateIV();

                return new AesConfig
                {
                    IsAesEnabled = true,
                    KeySize = keySize,
                    Key = aes.Key,
                    IV = aes.IV
                };
            }
        }

        public static AesConfig FromSharedSecret(string secret, AesKeySize keySize = AesKeySize.Aes256)
        {
            return new AesConfig
            {
                IsAesEnabled = true,
                SharedSecret = secret,
                KeySize = keySize
            };
        }

        public void Dispose()
        {
            if (!_disposed)
            {
                _encryptor?.Dispose();
                _decryptor?.Dispose();
                _aes?.Dispose();
                _disposed = true;
            }
        }
    }
}