// Любая реализация хранения slave-ключей, например в SharedPreferences
interface KeyStorage {
    operator fun get(name: String): ByteArray?
    operator fun set(name: String, key: ByteArray?)
}
 
class AesEncryptionManager(
    context: Context,
    initExecutor: Executor,
    exceptionHandler: (Exception) -> Unit,
    private val keyStorage: KeyStorage,
) : EncryptionManager {
 
    private val initLock = ReentrantReadWriteLock()
 
    private val masterKeyValidityStartDate: Date
    private val masterKeyValidityEndDate: Date
 
    private var initLatch = CountDownLatch(1)
 
    private lateinit var keyStore: KeyStore
    private lateinit var aesCipher: Cipher
    private val cipherLock: ReentrantLock = ReentrantLock()
 
    init {
        val calendar = Calendar.getInstance()
        masterKeyValidityStartDate = calendar.time
        calendar.add(Calendar.YEAR, 30)
        masterKeyValidityEndDate = calendar.time
        initExecutor.execute { init(exceptionHandler, masterKeyCreationCallback) }
    }
 
    @WorkerThread
    @Throws(EncryptionException::class)
    fun init(exceptionHandler: (Exception) -> Unit) = initLock.write {
        if (initLatch.count == 0L) {
            return
        }
 
        try {
            keyStore = KeyStore.getInstance("AndroidKeyStore")
            keyStore.load(null)
            aesCipher = Cipher.getInstance(AES_CIPHER_SUIT)
 
            if (!hasMasterKey()) {
                createMasterKey()
                masterKeyCreationCallback()
            }
        } catch (e: Exception) {
            exceptionHandler.invoke(EncryptionException("Failed to run init", e))
        } finally {
            initLatch.countDown()
        }
    }
 
    override fun encrypt(keyAlias: String, data: ByteArray): EncryptionManager.EncryptedData? {
        initLock.read {
            checkStateOrThrow()
        }
 
        val key = getKey(keyAlias) ?: createKey(keyAlias)
        return encryptWithKey(key, data)
    }
 
    override fun decrypt(keyAlias: String, data: EncryptionManager.EncryptedData): ByteArray? {
        initLock.read {
            checkStateOrThrow()
        }
 
        val key = getKey(keyAlias) ?: throw EncryptionException("No key with alias $keyAlias")
        return decryptWithKey(key, data)
    }
 
    override fun removeKey(keyAlias: String) {
        keyStorage[keyAlias] = null
    }
 
    override fun waitForInitialize(maxTimeMs: Long): Boolean {
        return initLatch.await(maxTimeMs, TimeUnit.MILLISECONDS)
    }
 
    // -- Master key operations
 
    private fun checkStateOrThrow() {
        if (initLatch.count > 0L) {
            throw EncryptionException("Manager is not initialized")
        }
 
        if (!hasMasterKey()) {
            throw EncryptionException("Cannot perform operations without master key")
        }
    }
 
    private fun createMasterKey() {
        try {
            KeyPairGenerator
                .getInstance("RSA", "AndroidKeyStore")
                .apply {
                    initialize(createSpec())
                    generateKeyPair()
                }
        } catch (e: Exception) {
            throw EncryptionException("Failed to generate master key", e)
        }
    }
 
    private fun hasMasterKey(): Boolean {
        return try {
            keyStore.getKey(MASTER_KEY_ALIAS, null) != null
        } catch (e: Exception) {
            L.w(e, "Failed to retrieve master key")
            false
        }
    }
 
    private fun encryptWithMasterKey(data: ByteArray): ByteArray {
        return try {
            val cipher = Cipher.getInstance(RSA_KEY_SUIT)
            val key = keyStore.getCertificate(MASTER_KEY_ALIAS).publicKey
            cipher.init(Cipher.ENCRYPT_MODE, key)
            cipher.doFinal(data)
        } catch (e: Exception) {
            throw EncryptionException("Failed to encrypt with master key", e)
        }
    }
 
    private fun decryptWithMasterKey(data: ByteArray): ByteArray {
        return try {
            val cipher = Cipher.getInstance(RSA_KEY_SUIT)
            val key = keyStore.getKey(MASTER_KEY_ALIAS, null)
            cipher.init(Cipher.DECRYPT_MODE, key)
            cipher.doFinal(data)
        } catch (e: Exception) {
            throw EncryptionException("Failed to decrypt with master key", e)
        }
    }
 
// -- AES key operations
 
    private fun getKey(name: String): Key? {
        val encryptedKey = keyStorage[name]
        if (encryptedKey == null) {
            L.i("No key with alias $name")
            return null
        }
 
        return Key(decryptWithMasterKey(encryptedKey))
    }
 
    private fun createKey(name: String): Key {
        val key = UUID.randomUUID().toString()
            .lowercase()
            .replace("-", "")
            .toCharArray()
 
        val salt = UUID.randomUUID().toByteArray()
 
        val spec = PBEKeySpec(key, salt, PBKDF2_ITER_COUNT, AES_KEY_SIZE)
 
        val skf = SecretKeyFactory.getInstance(PBKDF2_ALGORITHM)
        val generatedKey = try {
            skf.generateSecret(spec).encoded
        } catch (e: Exception) {
            throw EncryptionException("Failed to generate key", e)
        }
 
        val encryptedKey = encryptWithMasterKey(generatedKey)
        keyStorage[name] = encryptedKey
 
        return Key(generatedKey)
    }
 
    private fun encryptWithKey(key: Key, data: ByteArray): EncryptionManager.EncryptedData {
        return try {
            val keySpec = SecretKeySpec(key.encodedKey, AES_KEY_SPEC)
            cipherLock.withLock {
                aesCipher.init(Cipher.ENCRYPT_MODE, keySpec)
                val encrypted = aesCipher.doFinal(data)
                EncryptionManager.EncryptedData(encrypted, aesCipher.iv)
            }
        } catch (e: Exception) {
            throw EncryptionException("Failed to encrypt with raw aes key", e)
        }
    }
 
    private fun decryptWithKey(key: Key, data: EncryptionManager.EncryptedData): ByteArray {
        return try {
            cipherLock.withLock {
                val keySpec = SecretKeySpec(key.encodedKey, AES_KEY_SPEC)
                aesCipher.init(Cipher.DECRYPT_MODE, keySpec, IvParameterSpec(data.initVector))
                aesCipher.doFinal(data.data)
            }
        } catch (e: Exception) {
            throw EncryptionException("Failed to decrypt with aes key", e)
        }
    }
 
    private fun createSpec(): AlgorithmParameterSpec {
        return KeyGenParameterSpec.Builder(MASTER_KEY_ALIAS, KeyProperties.PURPOSE_ENCRYPT or KeyProperties.PURPOSE_DECRYPT)
            .setKeySize(KEY_SIZE)
            .setEncryptionPaddings(KeyProperties.ENCRYPTION_PADDING_RSA_PKCS1)
            .setAlgorithmParameterSpec(RSAKeyGenParameterSpec(KEY_SIZE, RSAKeyGenParameterSpec.F4))
            // key validity start & end explicitly not set as they caused KeyNotYetValidException @ api 27
            .setCertificateSubject(X500Principal("CN=$MASTER_KEY_ALIAS"))
            .setCertificateSerialNumber(BigInteger.valueOf(abs(MASTER_KEY_ALIAS.hashCode()).toLong()))
            .build()
    }
 
    companion object {
        private const val MASTER_KEY_ALIAS = "ALIAS_MASTER_KEY"
 
        private const val KEY_SIZE = 2048 // bits
        private const val AES_KEY_SIZE = 256
        private const val PBKDF2_ITER_COUNT = 10000
 
        private const val AES_CIPHER_SUIT = "AES/CBC/PKCS7Padding"
        private const val AES_KEY_SPEC = "AES"
        private const val RSA_KEY_SUIT = "RSA/NONE/PKCS1Padding"
        private const val PBKDF2_ALGORITHM = "PBKDF2WithHmacSHA1"
    }
}

class EncryptedFileWrapper(
    context: Context,
    fileName: String,
    private val sessionStatHelper: SessionStatHelper
) {
    private val lock = Any()
    
    private val encryptedFileLink by lazy(mode = LazyThreadSafetyMode.NONE) {
        File(context.filesDir, "encrypted_$fileName")
    }
    private val encryptedFile: EncryptedFile by lazy(mode = LazyThreadSafetyMode.NONE) {
        try {
            EncryptedFile.Builder(
                encryptedFileLink,
                context,
                MasterKeys.getOrCreate(MasterKeys.AES256_GCM_SPEC),
                EncryptedFile.FileEncryptionScheme.AES256_GCM_HKDF_4KB
            ).build()
        } catch (th: Throwable) {
            L.e(th, "create_encrypted_file")
            sessionStatHelper.sendError(mapOf("action" to "create_encrypted_file", "stacktrace" to th.stackTraceToString()))
            throw th
        }
    }

    @Throws(GeneralSecurityException::class, IOException::class, KeyStoreException::class)
    fun writeEncryptedUnsafe(value: String) = synchronized(lock) {
        if (encryptedFileLink.exists()) {
            encryptedFileLink.delete()
        }

        val bytes = value.toByteArray(StandardCharsets.UTF_8)
        encryptedFile.openFileOutput().use {
            it.write(bytes)
            it.flush()
        }
    }

    @Throws(GeneralSecurityException::class, IOException::class, KeyStoreException::class)
    fun readEncryptedUnsafe(): String? = synchronized(lock) { // blocking reading
        if (!encryptedFileLink.exists()) {
            return null
        }

        if (encryptedFileLink.length() == 0L) {
            encryptedFileLink.delete()
            return null
        }

        val byteArrayOutputStream = ByteArrayOutputStream()
        var nextByte: Int
        encryptedFile.openFileInput().use { encryptedStream ->
            nextByte = encryptedStream.read()
            while (nextByte != -1) {
                byteArrayOutputStream.write(nextByte)
                nextByte = encryptedStream.read()
            }
        }

        return byteArrayOutputStream.use { String(it.toByteArray(), StandardCharsets.UTF_8) }
    }

    @Throws(SecurityException::class)
    fun clear(): Boolean = synchronized(lock) {
        encryptedFileLink.delete()
    }
}

private val encryptedFile: EncryptedFile by lazy(mode = LazyThreadSafetyMode.NONE) {
    // Single-creation problem
    // Inspired by https://github.com/signalapp/Signal-Android/commit/f1f505d41c0d164c10226290d2af6f01b4461ae5
    try {
        createEncryptedFileForce(context)
    } catch (th: Throwable) {
        L.e(th, "create_encrypted_file_1")
        sessionStatHelper.sendError(mapOf("action" to create_encrypted_file_1", "stacktrace" to th.partStackTrace()))
 
        try {
            createEncryptedFileForce(context)
        } catch (th: Throwable) {
            L.e(th, "create_encrypted_file_1")
            sessionStatHelper.sendError(mapOf("action" to "EF_create_encrypted_file_2", "stacktrace" to th.partStackTrace()))
            throw th
        }
    }
}

class SafeEncryptedPreferences(
    context: Context,
    fileName: String
) : SharedPreferences {
    private val encrypted: SharedPreferences by lazy { EncryptedPreferencesHelper.createEncrypted(context, fileName, plain) }
    private val plain: SharedPreferences by lazy {
        context.getSharedPreferences("plain_$fileName", MODE_PRIVATE)
    }
 
    override fun contains(key: String?): Boolean {
        return encrypted.safeContains(key) || plain.contains(key)
    }
 
    override fun getBoolean(key: String?, defValue: Boolean): Boolean {
        return get(key, defValue, SharedPreferences::getBoolean)
    }
 
    override fun getInt(key: String?, defValue: Int): Int {
        return get(key, defValue, SharedPreferences::getInt)
    }
 
    override fun getLong(key: String?, defValue: Long): Long {
        return get(key, defValue, SharedPreferences::getLong)
    }
 
    override fun getFloat(key: String?, defValue: Float): Float {
        return get(key, defValue, SharedPreferences::getFloat)
    }
 
    override fun getString(key: String?, defValue: String?): String? {
        return get(key, defValue, SharedPreferences::getString)
    }
 
    override fun getStringSet(key: String?, defValues: MutableSet<String>?): MutableSet<String>? {
        return get(key, defValues, SharedPreferences::getStringSet)
    }
 
    override fun getAll(): MutableMap<String, *> {
        val allEncrypted = encrypted.safeAll()
        val allPlain = plain.all
        val all = HashMap<String, Any?>(allEncrypted.size + allEncrypted.size)
        all.putAll(allPlain)
        all.putAll(allEncrypted)
        return all
    }
 
    override fun edit(): SharedPreferences.Editor {
        return Editor(encrypted.edit(), plain.edit())
    }
 
    override fun registerOnSharedPreferenceChangeListener(
        listener: SharedPreferences.OnSharedPreferenceChangeListener?
    ) {
        encrypted.registerOnSharedPreferenceChangeListener(listener)
        plain.registerOnSharedPreferenceChangeListener(listener)
    }
 
    override fun unregisterOnSharedPreferenceChangeListener(
        listener: SharedPreferences.OnSharedPreferenceChangeListener?
    ) {
        encrypted.unregisterOnSharedPreferenceChangeListener(listener)
        plain.unregisterOnSharedPreferenceChangeListener(listener)
    }
 
    private inline fun <T> get(
        key: String?,
        defValue: T,
        getter: SharedPreferences.(key: String?, defValue: T) -> T
    ): T {
        return if (encrypted.safeContains(key)) {
            try {
                encrypted.getter(key, defValue)
            } catch (e: Exception) {
                plain.getter(key, defValue)
            }
        } else {
            plain.getter(key, defValue)
        }
    }
 
    private class Editor(
        private val encryptedEditor: SharedPreferences.Editor,
        private val plainEditor: SharedPreferences.Editor
    ) : SharedPreferences.Editor {
 
        private val clearRequested = AtomicBoolean(false)
 
        override fun remove(key: String?) = apply {
            encryptedEditor.safeRemove(key)
            plainEditor.remove(key)
        }
 
        override fun clear() = apply {
            clearRequested.set(true)
            encryptedEditor.safeClear()
            plainEditor.clear()
        }
 
        override fun putLong(key: String?, value: Long) = apply {
            put(key, value, SharedPreferences.Editor::putLong)
        }
 
        override fun putInt(key: String?, value: Int) = apply {
            put(key, value, SharedPreferences.Editor::putInt)
        }
 
        override fun putBoolean(key: String?, value: Boolean) = apply {
            put(key, value, SharedPreferences.Editor::putBoolean)
        }
 
        override fun putStringSet(key: String?, values: MutableSet<String>?) = apply {
            put(key, values, SharedPreferences.Editor::putStringSet)
        }
 
        override fun putFloat(key: String?, value: Float) = apply {
            put(key, value, SharedPreferences.Editor::putFloat)
        }
 
        override fun putString(key: String?, value: String?) = apply {
            put(key, value, SharedPreferences.Editor::putString)
        }
 
        override fun commit(): Boolean {
            return encryptedEditor.safeCommit() && plainEditor.commit()
        }
 
        override fun apply() {
            if (clearRequested.getAndSet(false)) {
                encryptedEditor.safeCommit()
            } else {
                encryptedEditor.safeApply()
            }
            plainEditor.apply()
        }
 
        private inline fun <T> put(
            key: String?,
            value: T,
            putter: SharedPreferences.Editor.(key: String?, value: T) -> Any
        ) = apply {
            try {
                encryptedEditor.putter(key, value)
            } catch (e: Exception) {
                plainEditor.putter(key, value)
            }
        }
    }
 
    companion object {
        fun SharedPreferences.safeContains(key: String?): Boolean {
            return try {
                contains(key)
            } catch (e: Exception) {
                false
            }
        }
 
        fun SharedPreferences.safeAll(): Map<String, *> {
            return try {
                all
            } catch (e: Exception) {
                emptyMap<String, Any?>()
            }
        }
 
        fun SharedPreferences.Editor.safeRemove(key: String?): SharedPreferences.Editor {
            return try {
                remove(key)
            } catch (e: Exception) {
                this
            }
        }
 
        fun SharedPreferences.Editor.safeClear(): SharedPreferences.Editor {
            return try {
                clear()
            } catch (e: Exception) {
                this
            }
        }
 
        fun SharedPreferences.Editor.safeCommit(): Boolean {
            return try {
                commit()
            } catch (e: Exception) {
                return false
            }
        }
 
        fun SharedPreferences.Editor.safeApply() {
            return try {
                apply()
            } catch (ignored: Exception) {
            }
        }
    }
}