Imports System.IO
Imports System.Security.Cryptography
Imports System.Text.Encoding
Imports DigitalData.Modules.Logging

''' <summary>
''' https://stackoverflow.com/questions/10168240/encrypting-decrypting-a-string-in-c-sharp
''' </summary>
Public Class Encryption
    ' This constant is used to determine the keysize of the encryption algorithm in bits.
    ' We divide this by 8 within the code below to get the equivalent number of bytes.
    Private Const KEY_SIZE As Integer = 256
    ' This constant determines the number of iterations for the password bytes generation function.
    Private Const DERIVATION_ITERATIONS As Integer = 1000
    Private Const BLOCK_SIZE As Integer = 256

    Private _paddingMode As PaddingMode = PaddingMode.Zeros
    Private _cipherMode As CipherMode = CipherMode.CBC

    Private ReadOnly _password As String
    Private _logger As Logger

    Public Sub New(LogConfig As LogConfig, Password As String)
        _logger = LogConfig.GetLogger()

        If IsNothing(Password) Then
            Throw New ArgumentNullException("Password")
        End If

        _password = Password
    End Sub

    Public Async Function EncryptAsync(PlainTextBytes As Byte()) As Task(Of Byte())
        Return Await Task.Run(Function() As Byte()
                                  Return Encrypt(PlainTextBytes)
                              End Function)
    End Function

    Public Function Encrypt(PlainText As String) As String
        Try
            Dim oBytes As Byte() = UTF8.GetBytes(PlainText)
            Dim oEncrypted As Byte() = Encrypt(oBytes)
            Return UTF8.GetString(oEncrypted)
        Catch ex As Exception
            _logger.Error(ex)
            Throw ex
        End Try
    End Function

    Public Function Encrypt(PlainTextBytes As Byte()) As Byte()
        Try
            ' Salt and IV is randomly generated each time, but is preprended to encrypted cipher text
            ' so that the same Salt and IV values can be used when decrypting.  
            Dim oSaltStringBytes = Generate256BitsOfRandomEntropy()
            Dim oIvStringBytes = Generate256BitsOfRandomEntropy()
            Using oPassword = New Rfc2898DeriveBytes(_password, oSaltStringBytes, DERIVATION_ITERATIONS)
                Dim oKeyBytes = oPassword.GetBytes(KEY_SIZE / 8)
                Using oSymmetricKey = New RijndaelManaged()
                    oSymmetricKey.BlockSize = BLOCK_SIZE
                    oSymmetricKey.Mode = _cipherMode
                    oSymmetricKey.Padding = _paddingMode

                    Using oEncryptor = oSymmetricKey.CreateEncryptor(oKeyBytes, oIvStringBytes)
                        Using oMemoryStream = New MemoryStream()
                            Using oCryptoStream = New CryptoStream(oMemoryStream, oEncryptor, CryptoStreamMode.Write)
                                oCryptoStream.Write(PlainTextBytes, 0, PlainTextBytes.Length)
                                oCryptoStream.FlushFinalBlock()
                                ' Create the final bytes as a concatenation of the random salt bytes, the random iv bytes and the cipher bytes.
                                Dim oCipherTextBytes = oSaltStringBytes
                                oCipherTextBytes = oCipherTextBytes.Concat(oIvStringBytes).ToArray()
                                oCipherTextBytes = oCipherTextBytes.Concat(oMemoryStream.ToArray()).ToArray()
                                oMemoryStream.Close()
                                oCryptoStream.Close()
                                Return oCipherTextBytes
                            End Using
                        End Using
                    End Using
                End Using
            End Using
        Catch ex As Exception
            _logger.Error(ex)
            Throw ex
        End Try
    End Function

    Public Async Function DecryptAsync(CipherTextBytesWithSaltAndIv As Byte()) As Task(Of Byte())
        Return Await Task.Run(Function() As Byte()
                                  Return Decrypt(CipherTextBytesWithSaltAndIv)
                              End Function)
    End Function

    Public Function Decrypt(CipherTextPlainWithSaltAndIv As String) As String
        Try
            Dim oBytes As Byte() = UTF8.GetBytes(CipherTextPlainWithSaltAndIv)
            Dim oDecrypted As Byte() = Decrypt(oBytes)
            Return UTF8.GetString(oDecrypted)
        Catch ex As Exception
            _logger.Error(ex)
            Throw ex
        End Try
    End Function

    Public Function Decrypt(CipherTextBytesWithSaltAndIv As Byte()) As Byte()
        Try
            ' Get the complete stream of bytes that represent:
            ' [32 bytes of Salt] + [32 bytes of IV] + [n bytes of CipherText]
            ' Get the saltbytes by extracting the first 32 bytes from the supplied cipherText bytes.
            Dim oSaltStringBytes = CipherTextBytesWithSaltAndIv.Take(KEY_SIZE / 8).ToArray()
            ' Get the IV bytes by extracting the next 32 bytes from the supplied cipherText bytes.
            Dim oIvStringBytes = CipherTextBytesWithSaltAndIv.Skip(KEY_SIZE / 8).Take(KEY_SIZE / 8).ToArray()
            ' Get the actual cipher text bytes by removing the first 64 bytes from the cipherText string.
            Dim oCipherTextBytes = CipherTextBytesWithSaltAndIv.Skip((KEY_SIZE / 8) * 2).Take(CipherTextBytesWithSaltAndIv.Length - ((KEY_SIZE / 8) * 2)).ToArray()

            Using oPassword = New Rfc2898DeriveBytes(_password, oSaltStringBytes, DERIVATION_ITERATIONS)
                Dim oKeyBytes = oPassword.GetBytes(KEY_SIZE / 8)
                Using oSymmetricKey = New RijndaelManaged()
                    oSymmetricKey.BlockSize = BLOCK_SIZE
                    oSymmetricKey.Mode = _cipherMode
                    oSymmetricKey.Padding = _paddingMode
                    Using oDecryptor = oSymmetricKey.CreateDecryptor(oKeyBytes, oIvStringBytes)
                        Using oMemoryStream = New MemoryStream(oCipherTextBytes)
                            Using oCryptoStream = New CryptoStream(oMemoryStream, oDecryptor, CryptoStreamMode.Read)
                                Dim oPlainTextBytes = New Byte(oCipherTextBytes.Length - 1) {}
                                Dim oDecryptedByteCount = oCryptoStream.Read(oPlainTextBytes, 0, oPlainTextBytes.Length)
                                oMemoryStream.Close()
                                oCryptoStream.Close()
                                Return oPlainTextBytes
                            End Using
                        End Using
                    End Using
                End Using
            End Using
        Catch ex As Exception
            _logger.Error(ex)
            Throw ex
        End Try
    End Function

    Private Shared Function Generate256BitsOfRandomEntropy() As Byte()
        Dim oRandomBytes = New Byte(31) {}
        ' 32 Bytes will give us 256 bits.
        Using oRNGCsp = New RNGCryptoServiceProvider()
            ' Fill the array with cryptographically secure random bytes.
            oRNGCsp.GetBytes(oRandomBytes)
        End Using
        Return oRandomBytes
    End Function
End Class