/*
 * Copyright 2014-2022 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
 */

package io.ktor.server.sessions

import io.ktor.util.*
import org.slf4j.*
import java.security.*
import javax.crypto.*
import javax.crypto.spec.*

/**
 * A transformer used to sign and encrypt/decrypt session data.
 * This transformer works as follows:
 * - encrypts/decrypts data using [encryptAlgorithm] and [encryptionKeySpec]
 * - includes an authenticated MAC (Message Authentication Code) hash with [signAlgorithm] and [signKeySpec]
 * - includes an IV (Initialization Vector) that is generated by an [ivGenerator] by default secure random bytes
 *
 * By default, it uses AES for encryption and HmacSHA256 for authentication.
 *
 * You have to provide keys of compatible sizes: 16, 24 and 32 for AES encryption.
 * For HmacSHA256 it is recommended a key of 32 bytes.
 *
 * @see [Sessions]
 *
 * @property encryptionKeySpec is a secret key that is used for encryption
 * @property signKeySpec is a secret key that is used for signing
 * @property ivGenerator is a function that generates input vectors
 * @property encryptAlgorithm is an encryption algorithm name
 * @property signAlgorithm is a signing algorithm name
 * @property backwardCompatibleRead before Ktor 3.0.0, MAC was calculated over decrypted data.
 * Set to true to support old clients.
 */
public class SessionTransportTransformerEncrypt(
    public val encryptionKeySpec: SecretKeySpec,
    public val signKeySpec: SecretKeySpec,
    public val ivGenerator: (size: Int) -> ByteArray =
        { size -> ByteArray(size).apply { SecureRandom().nextBytes(this) } },
    public val encryptAlgorithm: String = encryptionKeySpec.algorithm,
    public val signAlgorithm: String = signKeySpec.algorithm,
    private val backwardCompatibleRead: Boolean = false,
) : SessionTransportTransformer {
    public companion object {
        private val log = LoggerFactory.getLogger(SessionTransportTransformerEncrypt::class.qualifiedName)
    }

    private val charset = Charsets.UTF_8

    /**
     * A size of the key used to encrypt session data.
     */
    public val encryptionKeySize: Int get() = encryptionKeySpec.encoded.size

    // Check that input keys are right
    init {
        encrypt(ivGenerator(encryptionKeySize), byteArrayOf())
        mac(byteArrayOf())
    }

    public constructor(
        encryptionKey: ByteArray,
        signKey: ByteArray,
        ivGenerator: (size: Int) -> ByteArray = { size -> ByteArray(size).apply { SecureRandom().nextBytes(this) } },
        encryptAlgorithm: String = "AES",
        signAlgorithm: String = "HmacSHA256",
        backwardCompatibleRead: Boolean = false,
    ) : this(
        encryptionKeySpec = SecretKeySpec(encryptionKey, encryptAlgorithm),
        signKeySpec = SecretKeySpec(signKey, signAlgorithm),
        ivGenerator = ivGenerator,
        backwardCompatibleRead = backwardCompatibleRead
    )

    override fun transformRead(transportValue: String): String? {
        try {
            val encryptedAndMac = transportValue.substringAfterLast('/', "")
            val macHex = encryptedAndMac.substringAfterLast(':', "")
            val encrypted = hex(encryptedAndMac.substringBeforeLast(':'))
            val macCheck = hex(mac(encrypted)) == macHex
            if (!macCheck && !backwardCompatibleRead) {
                return null
            }

            val iv = hex(transportValue.substringBeforeLast('/'))
            val decrypted = decrypt(iv, encrypted)

            if (!macCheck && hex(mac(decrypted)) != macHex) {
                return null
            }

            return decrypted.toString(charset)
        } catch (e: Throwable) {
            // NumberFormatException // Invalid hex
            // InvalidAlgorithmParameterException // Invalid data
            if (log.isDebugEnabled) {
                log.debug(e.toString())
            }
            return null
        }
    }

    override fun transformWrite(transportValue: String): String {
        val iv = ivGenerator(encryptionKeySize)
        val decrypted = transportValue.toByteArray(charset)
        val encrypted = encrypt(iv, decrypted)
        val mac = mac(encrypted)
        return "${hex(iv)}/${hex(encrypted)}:${hex(mac)}"
    }

    private fun encrypt(initVector: ByteArray, decrypted: ByteArray): ByteArray {
        return encryptDecrypt(Cipher.ENCRYPT_MODE, initVector, decrypted)
    }

    private fun decrypt(initVector: ByteArray, encrypted: ByteArray): ByteArray {
        return encryptDecrypt(Cipher.DECRYPT_MODE, initVector, encrypted)
    }

    private fun encryptDecrypt(mode: Int, initVector: ByteArray, input: ByteArray): ByteArray {
        val iv = IvParameterSpec(initVector)
        val cipher = Cipher.getInstance("$encryptAlgorithm/CBC/PKCS5PADDING")
        cipher.init(mode, encryptionKeySpec, iv)
        return cipher.doFinal(input)
    }

    private fun mac(value: ByteArray): ByteArray = Mac.getInstance(signAlgorithm).run {
        init(signKeySpec)
        doFinal(value)
    }
}
