Key Wrapping example

KeyWrapper.kt

import android.security.keystore.KeyProperties
import android.util.Base64
import java.security.Key
import java.security.SecureRandom
import javax.crypto.Cipher
import javax.crypto.spec.IvParameterSpec
import javax.crypto.spec.SecretKeySpec
import kotlin.jvm.Throws

// Source - https://medium.com/@hayk.mkrtchyan8998/shedding-light-on-android-encryption-android-crypto-api-part-2-cipher-147ff4411e1d

interface KeyWrapper {
    /**
     * Get the key that is saved in a storage or create a new one if it doesn't exist.
     * Note: Wrapping logic is already handled by this function.
     * @return - The key
     */
    @Throws(Exception::class)
    fun getOrCreateKey(): Key
}

class KeyWrapperImpl : KeyWrapper {

    /**
     * The key length should be 16, 24 or 32 characters long
     * - 16 -> AES-128
     * - 24 -> AES-192
     * - 32 -> AES-256
     * @see <a href="https://en.wikipedia.org/wiki/Advanced_Encryption_Standard">AES Algorithm</a>
     */
    private val baseKeyMaterial = "keep_this_key_sc".toByteArray(Charsets.UTF_8)

    /**
     * Out initial key, using which we will wrap and unwrap other key.
     * wrapKey and unwrapKey functions accept initialKey as an argument.
     * It's better to keep it outside of this class to make this class more generic and reusable.
     * Note: You can also use KeyGenerator to generate a key. It's out of this lesson scope.
     * I will tell about it in the next part.
     */
    private val initialKey = SecretKeySpec(baseKeyMaterial, AES_ALGORITHM)

    /**
     * Function that wraps the key with another key
     * @param initialKey - The key that will be used to wrap the other key
     * @param keyToBeWrapped - The key that will be wrapped
     * @return - The encrypted key data
     */
    @Throws(Exception::class)
    private fun wrapKey(initialKey: Key, keyToBeWrapped: Key): ByteArray {
        val cipher = Cipher.getInstance(TRANSFORMATION)
        cipher.init(Cipher.WRAP_MODE, initialKey)
        val encryptedKey = cipher.wrap(keyToBeWrapped)

        val iv = cipher.iv
        val encryptedDataWithIV = ByteArray(iv.size + encryptedKey.size)
        System.arraycopy(iv, 0, encryptedDataWithIV, 0, iv.size)
        System.arraycopy(encryptedKey, 0, encryptedDataWithIV, iv.size, encryptedKey.size)

        return encryptedDataWithIV
    }

    /**
     * Function that unwraps the key. Use the same initial key that was used to wrap the key.
     * @param initialKey - The key that will be used to wrap the other key
     * @param wrappedKey - The encrypted key data
     * @return - The unwrapped key
     */
    @Throws(Exception::class)
    private fun unwrapKey(initialKey: Key, wrappedKey: ByteArray): Key {
        val cipher = Cipher.getInstance(TRANSFORMATION)
        val iv = wrappedKey.copyOfRange(0, cipher.blockSize)
        val encryptedData = wrappedKey.copyOfRange(cipher.blockSize, wrappedKey.size)
        cipher.init(Cipher.UNWRAP_MODE, initialKey, IvParameterSpec(iv))
        return cipher.unwrap(encryptedData, AES_ALGORITHM, Cipher.SECRET_KEY)
    }

    /**
     * Creates a new key. You can also use KeyGenerator to generate a key.
     * @param keyMaterial - The material for the key
     * @return - The newly created key
     */
    @Throws(IllegalArgumentException::class)
    private fun createKey(keyMaterial: ByteArray): Key {
        return SecretKeySpec(keyMaterial, AES_ALGORITHM)
    }

    /**
     * Generates a random key material. This adds an extra layer of security as the material is random.
     */
    private fun generateRandomKeyMaterial(): ByteArray {
        val random = SecureRandom()
        val key = ByteArray(16)
        random.nextBytes(key)
        return key
    }

    @Throws(Exception::class)
    override fun getOrCreateKey(): Key {
        
        //Retrieve the key from where you have persisted it. Currently setting to null for demonstration purposes
        val existingKeyFromStorage: String? = null

        return existingKeyFromStorage?.let {
            // If the key is found in the a storage, we unwrap and return it
            unwrapKey(
                initialKey = initialKey,
                wrappedKey = Base64.decode(it, Base64.DEFAULT),
            )
        } ?: run {
            /**
             * Key does not exist in local storage. Maybe it's our first run or the user cleared the app data
             * Create a new key, wrap it, save in local storage and return it
             * NOTE - If you have some encrypted data with your generated key, but you have lost it,
             * you won't be able to decrypt it. So make sure to save it somewhere.
             */
            val key = createKey(generateRandomKeyMaterial())
            val encryptedKey = wrapKey(
                initialKey = initialKey,
                keyToBeWrapped = key,
            )

            // Persist this somewhere
            Base64.encodeToString(encryptedKey, Base64.DEFAULT)

            key
        }
    }

    companion object {
        private const val AES_ALGORITHM = KeyProperties.KEY_ALGORITHM_AES
        private const val BLOCK_MODE = KeyProperties.BLOCK_MODE_CBC
        private const val PADDING = KeyProperties.ENCRYPTION_PADDING_PKCS7
        private const val TRANSFORMATION = "$AES_ALGORITHM/$BLOCK_MODE/$PADDING"
    }
}