SecureStorageManager.kt

import android.content.Context
import android.os.Build
import android.util.Base64
import androidx.annotation.RequiresApi
import java.security.SecureRandom
import javax.crypto.Cipher
import javax.crypto.KeyGenerator
import javax.crypto.Mac
import javax.crypto.SecretKey
import javax.crypto.spec.GCMParameterSpec
import javax.crypto.spec.IvParameterSpec
import javax.crypto.spec.SecretKeySpec

class SecureStorageManager(context: Context) {

    companion object {
        private const val AES_TRANSFORMATION = "AES/GCM/NoPadding"
        private const val AES_LEGACY_TRANSFORMATION = "AES/CBC/PKCS5Padding"
        private const val HMAC_ALGORITHM = "HmacSHA256"
        private const val PREFS_NAME = "secure_storage"
        private const val ENCRYPTED_KEY_ALIAS = "encrypted_key"
    }

    internal val sharedPreferences = context.getSharedPreferences(PREFS_NAME, Context.MODE_PRIVATE)

    // Save encrypted data
    fun encryptAndSave(key: String, data: String) {
        val (iv, encryptedData) = if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.M) {
            encryptWithKeystore(data)
        } else {
            encryptWithAES(data)
        }

        val hmac = generateHMAC(encryptedData)

        sharedPreferences.edit()
            .putString("$key.data", Base64.encodeToString(encryptedData, Base64.DEFAULT))
            .putString("$key.iv", Base64.encodeToString(iv, Base64.DEFAULT))
            .putString("$key.hmac", Base64.encodeToString(hmac, Base64.DEFAULT))
            .apply()
    }

    // Retrieve and decrypt data
    fun decrypt(key: String): String? {
        val encryptedData = sharedPreferences.getString("$key.data", null)?.let {
            Base64.decode(it, Base64.DEFAULT)
        } ?: return null

        val iv = sharedPreferences.getString("$key.iv", null)?.let {
            Base64.decode(it, Base64.DEFAULT)
        } ?: return null

        val storedHmac = sharedPreferences.getString("$key.hmac", null)?.let {
            Base64.decode(it, Base64.DEFAULT)
        } ?: return null

        if (!verifyHMAC(encryptedData, storedHmac)) {
            throw SecurityException("HMAC verification failed")
        }

        return if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.M) {
            decryptWithKeystore(encryptedData, iv)
        } else {
            decryptWithAES(encryptedData, iv)
        }
    }

    // Encrypt with Keystore for Android >= 23
    @RequiresApi(Build.VERSION_CODES.M)
    private fun encryptWithKeystore(data: String): Pair<ByteArray, ByteArray> {
        val cipher = Cipher.getInstance(AES_TRANSFORMATION)
        val key = getOrCreateKeyForKeystore()
        cipher.init(Cipher.ENCRYPT_MODE, key)
        val iv = cipher.iv
        val encryptedData = cipher.doFinal(data.toByteArray())
        return Pair(iv, encryptedData)
    }

    // Decrypt with Keystore for Android >= 23
    @RequiresApi(Build.VERSION_CODES.M)
    private fun decryptWithKeystore(encryptedData: ByteArray, iv: ByteArray): String {
        val cipher = Cipher.getInstance(AES_TRANSFORMATION)
        val key = getOrCreateKeyForKeystore()
        cipher.init(Cipher.DECRYPT_MODE, key, GCMParameterSpec(128, iv))
        val decryptedBytes = cipher.doFinal(encryptedData)
        return String(decryptedBytes)
    }

    // Get or create a Keystore key
    @RequiresApi(Build.VERSION_CODES.M)
    private fun getOrCreateKeyForKeystore(): SecretKey {
        val keyStore = java.security.KeyStore.getInstance("AndroidKeyStore").apply { load(null) }
        val alias = "SecureStorageKey"
        if (!keyStore.containsAlias(alias)) {
            val keyGenerator = KeyGenerator.getInstance("AES", "AndroidKeyStore")
            keyGenerator.init(
                android.security.keystore.KeyGenParameterSpec.Builder(
                    alias,
                    android.security.keystore.KeyProperties.PURPOSE_ENCRYPT or android.security.keystore.KeyProperties.PURPOSE_DECRYPT
                )
                    .setBlockModes(android.security.keystore.KeyProperties.BLOCK_MODE_GCM)
                    .setEncryptionPaddings(android.security.keystore.KeyProperties.ENCRYPTION_PADDING_NONE)
                    .build()
            )
            keyGenerator.generateKey()
        }
        return keyStore.getKey(alias, null) as SecretKey
    }

    // Encrypt with AES for Android < 23
    private fun encryptWithAES(data: String): Pair<ByteArray, ByteArray> {
        val cipher = Cipher.getInstance(AES_LEGACY_TRANSFORMATION)
        val iv = ByteArray(16).apply { SecureRandom().nextBytes(this) }
        val key = getLegacyAESKey()
        cipher.init(Cipher.ENCRYPT_MODE, key, IvParameterSpec(iv))
        val encryptedData = cipher.doFinal(data.toByteArray())
        return Pair(iv, encryptedData)
    }

    // Decrypt with AES for Android < 23
    private fun decryptWithAES(encryptedData: ByteArray, iv: ByteArray): String {
        val cipher = Cipher.getInstance(AES_LEGACY_TRANSFORMATION)
        val key = getLegacyAESKey()
        cipher.init(Cipher.DECRYPT_MODE, key, IvParameterSpec(iv))
        val decryptedBytes = cipher.doFinal(encryptedData)
        return String(decryptedBytes)
    }

    // Generate or retrieve legacy AES key for Android < 23
    private fun getLegacyAESKey(): SecretKey {
        val existingKey = sharedPreferences.getString(ENCRYPTED_KEY_ALIAS, null)
        return if (existingKey != null) {
            val decodedKey = Base64.decode(existingKey, Base64.DEFAULT)
            SecretKeySpec(decodedKey, "AES")
        } else {
            val keyGenerator = KeyGenerator.getInstance("AES")
            keyGenerator.init(256)
            val newKey = keyGenerator.generateKey()
            sharedPreferences.edit()
                .putString(ENCRYPTED_KEY_ALIAS, Base64.encodeToString(newKey.encoded, Base64.DEFAULT))
                .apply()
            newKey
        }
    }

    // Generate HMAC for data
    private fun generateHMAC(data: ByteArray): ByteArray {
        val key = getLegacyAESKey()
        val mac = Mac.getInstance(HMAC_ALGORITHM)
        mac.init(SecretKeySpec(key.encoded, HMAC_ALGORITHM))
        return mac.doFinal(data)
    }

    // Verify HMAC for data
    private fun verifyHMAC(data: ByteArray, hmac: ByteArray): Boolean {
        val calculatedHMAC = generateHMAC(data)
        return calculatedHMAC.contentEquals(hmac)
    }
}