package name.blackcap.passman

import java.nio.ByteBuffer
import java.nio.CharBuffer
import java.nio.charset.Charset
import java.nio.charset.StandardCharsets
import java.security.SecureRandom
import javax.crypto.Cipher
import javax.crypto.SecretKey
import javax.crypto.SecretKeyFactory
import javax.crypto.spec.IvParameterSpec
import javax.crypto.spec.PBEKeySpec
import javax.crypto.spec.SecretKeySpec
import javax.security.auth.Destroyable

class Encryption(passwordIn: CharArray, saltIn: ByteArray) : Destroyable {
    private companion object {
        const val ITERATIONS = 390000
        const val IV_LENGTH = 16
        const val KEY_LENGTH = 256
        const val ENCRYPTION_ALGORITHM = "AES/CBC/PKCS5Padding"
        const val KEY_ALGORITHM = "AES"
        const val SECRET_KEY_FACTORY = "PBKDF2WithHmacSHA256"
        val CHARSET : Charset = StandardCharsets.UTF_8
        val ZERO_IV = ByteArray(IV_LENGTH).apply { clear() }
    }

    private val secretKey = getSecretKey(passwordIn, saltIn)
    private val secureRandom = SecureRandom()

    fun encrypt(plaintext: CharArray): ByteArray {
        val iv = ByteArray(IV_LENGTH).also { secureRandom.nextBytes(it) }
        return encrypt(plaintext, iv)
    }

    private fun encrypt(plaintext: CharArray, iv: ByteArray): ByteArray {
        val cipher = Cipher.getInstance(ENCRYPTION_ALGORITHM)
        cipher.init(Cipher.ENCRYPT_MODE, secretKey, ivFactory(iv))
        val inBuffer = CHARSET.encode(CharBuffer.wrap(plaintext))
        val outBuffer = ByteBuffer.allocate(cipher.getOutputSize(inBuffer.limit()) + IV_LENGTH)
        outBuffer.put(iv)
        cipher.doFinal(inBuffer, outBuffer)
        return outBuffer.array()
    }

    fun encryptFromString(plaintext: String): ByteArray = encrypt(plaintext.toCharArray())

    fun encryptFromString0(plaintext: String): ByteArray =
        encrypt(plaintext.toCharArray(), ZERO_IV)

    fun decrypt(ciphertext: ByteArray): CharArray {
        val cipher = Cipher.getInstance(ENCRYPTION_ALGORITHM)
        cipher.init(Cipher.DECRYPT_MODE, secretKey, ivFactory(ciphertext))
        val bytes = cipher.doFinal(ciphertext, IV_LENGTH, ciphertext.size - IV_LENGTH)
        val charBuffer = CHARSET.decode(ByteBuffer.wrap(bytes))
        bytes.clear()
        val ret = CharArray(charBuffer.limit())
        charBuffer.run {
            rewind()
            get(ret)
            zero()
        }
        return ret
    }

    fun decryptToString(ciphertext: ByteArray): String = String(decrypt(ciphertext))

    override fun destroy() {
        secretKey.destroy()
    }

    override fun isDestroyed(): Boolean {
        return secretKey.isDestroyed
    }

    protected fun finalize() {
        destroy()
    }

    private fun getSecretKey(password: CharArray, salt: ByteArray): SecretKey {
        val factory = SecretKeyFactory.getInstance(SECRET_KEY_FACTORY)
        val spec = PBEKeySpec(password, salt, ITERATIONS, KEY_LENGTH)
        return SecretKeySpec(factory.generateSecret(spec).encoded, KEY_ALGORITHM)
    }

    private fun ivFactory(ciphertext: ByteArray) = IvParameterSpec(ciphertext, 0, IV_LENGTH)
}

fun ByteArray.clear() = indices.forEach { this[it] = 0 }

fun CharArray.clear() = indices.forEach { this[it] = '\u0000' }

fun CharBuffer.zero() {
    clear()
    array().clear()
}