Imports System.IO
Imports System.Security.Cryptography

Namespace Security
    ''' <summary>
    ''' Provides secure AES file encryption and decryption using PBKDF2 (Rfc2898DeriveBytes) for key derivation.
    ''' File format:
    ''' [1 byte Version][4 bytes IterationCount (Int32, big-endian)][32 bytes Salt][Encrypted Payload]
    ''' </summary>
    Public NotInheritable Class SecureFileHandler
        Private Sub New()
        End Sub

        Private Const CURRENT_VERSION As Byte = 1
        Private Const SALT_LENGTH As Integer = 32
        Private Const KEY_SIZE_BYTES As Integer = 32 ' AES-256
        Private Const IV_SIZE_BYTES As Integer = 16 ' AES Block size (128 bit)
        Private Const DEFAULT_ITERATIONS As Integer = 100000
        Private Const BUFFER_SIZE As Integer = 81920 '80KB streaming buffer

        ''' <summary>
        ''' Encrypts the provided byte array and writes an encrypted file to the target path using streaming.
        ''' </summary>
        Public Shared Sub EncryptFileFromBytes(sourceData() As Byte, targetFilePath As String, password As String, Optional iterations As Integer = DEFAULT_ITERATIONS)
            If sourceData Is Nothing Then Throw New ArgumentNullException(NameOf(sourceData))
            If String.IsNullOrWhiteSpace(password) Then Throw New ArgumentNullException(NameOf(password))
            Using fsOut = New FileStream(targetFilePath, FileMode.Create, FileAccess.Write, FileShare.None)
                Dim salt = GenerateRandomBytes(SALT_LENGTH)

                ' Write header: Version, Iterations, Salt
                fsOut.WriteByte(CURRENT_VERSION)
                WriteInt32BigEndian(fsOut, iterations)
                fsOut.Write(salt, 0, salt.Length)

                Using keyDerivation = New Rfc2898DeriveBytes(password, salt, iterations)
                    Dim key = keyDerivation.GetBytes(KEY_SIZE_BYTES)
                    Dim iv = keyDerivation.GetBytes(IV_SIZE_BYTES)

                    Dim aesAlg As System.Security.Cryptography.Aes = System.Security.Cryptography.Aes.Create()
                    Try
                        aesAlg.KeySize = KEY_SIZE_BYTES * 8
                        aesAlg.BlockSize = IV_SIZE_BYTES * 8
                        aesAlg.Mode = CipherMode.CBC
                        aesAlg.Padding = PaddingMode.PKCS7
                        aesAlg.Key = key
                        aesAlg.IV = iv

                        Using crypto = aesAlg.CreateEncryptor()
                            Using cs = New CryptoStream(fsOut, crypto, CryptoStreamMode.Write)
                                Using msIn = New MemoryStream(sourceData, writable:=False)
                                    Dim buffer(BUFFER_SIZE - 1) As Byte
                                    Dim read As Integer
                                    Do
                                        read = msIn.Read(buffer, 0, buffer.Length)
                                        If read <= 0 Then Exit Do
                                        cs.Write(buffer, 0, read)
                                    Loop
                                End Using
                                cs.FlushFinalBlock()
                            End Using
                        End Using
                    Finally
                        aesAlg.Dispose()
                    End Try
                End Using
            End Using
        End Sub

        ''' <summary>
        ''' Decrypts the encrypted file and returns the plaintext bytes.
        ''' </summary>
        Public Shared Function DecryptFileToBytes(encryptedFilePath As String, password As String) As Byte()
            If String.IsNullOrWhiteSpace(encryptedFilePath) Then Throw New ArgumentNullException(NameOf(encryptedFilePath))
            If String.IsNullOrWhiteSpace(password) Then Throw New ArgumentNullException(NameOf(password))
            Using fsIn = New FileStream(encryptedFilePath, FileMode.Open, FileAccess.Read, FileShare.Read)
                Dim version = CByte(fsIn.ReadByte())
                If version <> CURRENT_VERSION Then Throw New InvalidDataException("Unsupported file version.")
                Dim iterations = ReadInt32BigEndian(fsIn)
                Dim salt = New Byte(SALT_LENGTH - 1) {}
                ReadExact(fsIn, salt, 0, salt.Length)

                Using keyDerivation = New Rfc2898DeriveBytes(password, salt, iterations)
                    Dim key = keyDerivation.GetBytes(KEY_SIZE_BYTES)
                    Dim iv = keyDerivation.GetBytes(IV_SIZE_BYTES)

                    Dim aesAlg As System.Security.Cryptography.Aes = System.Security.Cryptography.Aes.Create()
                    Try
                        aesAlg.KeySize = KEY_SIZE_BYTES * 8
                        aesAlg.BlockSize = IV_SIZE_BYTES * 8
                        aesAlg.Mode = CipherMode.CBC
                        aesAlg.Padding = PaddingMode.PKCS7
                        aesAlg.Key = key
                        aesAlg.IV = iv

                        Using crypto = aesAlg.CreateDecryptor()
                            Using cs = New CryptoStream(fsIn, crypto, CryptoStreamMode.Read)
                                Using msOut = New MemoryStream()
                                    Dim buffer(BUFFER_SIZE - 1) As Byte
                                    Dim read As Integer
                                    Do
                                        read = cs.Read(buffer, 0, buffer.Length)
                                        If read <= 0 Then Exit Do
                                        msOut.Write(buffer, 0, read)
                                    Loop
                                    Return msOut.ToArray()
                                End Using
                            End Using
                        End Using
                    Finally
                        aesAlg.Dispose()
                    End Try
                End Using
            End Using
        End Function

        Private Shared Function GenerateRandomBytes(length As Integer) As Byte()
            Dim data = New Byte(length - 1) {}
            Using rng = RandomNumberGenerator.Create()
                rng.GetBytes(data)
            End Using
            Return data
        End Function

        Private Shared Sub WriteInt32BigEndian(stream As Stream, value As Integer)
            Dim bytes = BitConverter.GetBytes(value)
            If BitConverter.IsLittleEndian Then Array.Reverse(bytes)
            stream.Write(bytes, 0, bytes.Length)
        End Sub

        Private Shared Function ReadInt32BigEndian(stream As Stream) As Integer
            Dim bytes = New Byte(3) {}
            ReadExact(stream, bytes, 0, 4)
            If BitConverter.IsLittleEndian Then Array.Reverse(bytes)
            Return BitConverter.ToInt32(bytes, 0)
        End Function

        Private Shared Sub ReadExact(stream As Stream, buffer As Byte(), offset As Integer, count As Integer)
            Dim totalRead As Integer = 0
            While totalRead < count
                Dim read = stream.Read(buffer, offset + totalRead, count - totalRead)
                If read <= 0 Then Throw New EndOfStreamException("Unexpected end of stream.")
                totalRead += read
            End While
        End Sub
    End Class
End Namespace