JWE RSA_OAEP_256 + A256GCM encryption / decryption

JWE.kt

import org.json.JSONObject
import java.net.URLDecoder
import java.net.URLEncoder
import java.security.AlgorithmParameters
import java.security.KeyPair
import java.security.KeyPairGenerator
import java.security.SecureRandom
import java.security.interfaces.RSAPrivateKey
import java.security.interfaces.RSAPublicKey
import java.security.spec.MGF1ParameterSpec
import java.util.Base64
import javax.crypto.Cipher
import javax.crypto.spec.GCMParameterSpec
import javax.crypto.spec.OAEPParameterSpec
import javax.crypto.spec.PSource
import javax.crypto.spec.SecretKeySpec

private const val IV_BIT_LENGTH = 96
private const val AUTH_TAG_BIT_LENGTH = 128
private const val AES_GCM_KEY_BIT_LENGTH = 256

fun main() {
    val keyPair = generateKeyPair()
    val publicKey = keyPair.public as RSAPublicKey
    val privateKey = keyPair.private as RSAPrivateKey

    val publicKeyModulus = publicKey.modulus
    val publicKeyExponent = publicKey.publicExponent

    println("n Modulus : " + publicKeyModulus.toByteArray().encodeUrlSafeBase64())
    println("e Modulus : " + publicKeyExponent.toByteArray().encodeUrlSafeBase64())

    val payload = Json {
        "first_name" to "John"
        "last_name" to "Doe"
    }.toString().toByteArray()

    val token = encode(publicKey, payload)
    println(token)
    val decoded = decode(privateKey, token)
    println(decoded)
}

private fun decode(privateKey: RSAPrivateKey, token: String): String {
    val splits = token.split(".")

    val header = splits[0]
    val encryptedCekKey = splits[1].decodeUrlSafeBase64()
    val iv = splits[2].decodeUrlSafeBase64()
    val cipherText = splits[3].decodeUrlSafeBase64()
    val authTag = splits[4].decodeUrlSafeBase64()

    val algParam = AlgorithmParameters.getInstance("OAEP")
    algParam.init(
        OAEPParameterSpec(
            "SHA-256",
            "MGF1",
            MGF1ParameterSpec.SHA256,
            PSource.PSpecified.DEFAULT
        )
    )
    val cekKeyBytes = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding").run {
        init(Cipher.DECRYPT_MODE, privateKey, algParam)
        doFinal(encryptedCekKey)
    }
    val cekKey = SecretKeySpec(cekKeyBytes, "AES")
    val gcmSpec = GCMParameterSpec(AUTH_TAG_BIT_LENGTH, iv)

    val cipherOutput = Cipher.getInstance("AES/GCM/NoPadding").run {
        init(Cipher.DECRYPT_MODE, cekKey, gcmSpec)
        updateAAD(header.toByteArray())
        doFinal(cipherText + authTag)
    }
    return String(cipherOutput)
}

private fun encode(publicKey: RSAPublicKey, payload: ByteArray): String {
    val header = Json {
        "enc" to "A256GCM"
        "alg" to "RSA-OAEP-256"
    }.toString().toByteArray().encodeUrlSafeBase64()

    val cekKey = SecretKeySpec(randomBytes(AES_GCM_KEY_BIT_LENGTH), "AES")
    val gcmSpec = GCMParameterSpec(AUTH_TAG_BIT_LENGTH, randomBytes(IV_BIT_LENGTH))

    val cipherOutput = Cipher.getInstance("AES/GCM/NoPadding").run {
        init(Cipher.ENCRYPT_MODE, cekKey, gcmSpec)
        updateAAD(header.toByteArray())
        doFinal(payload)
    }

    val tagPos: Int = cipherOutput.size - (AUTH_TAG_BIT_LENGTH / 8)
    val cipherTextBytes = subArray(cipherOutput, 0, tagPos)
    val authTagBytes = subArray(cipherOutput, tagPos, (AUTH_TAG_BIT_LENGTH / 8))

    val algParam = AlgorithmParameters.getInstance("OAEP")
    algParam.init(
        OAEPParameterSpec(
            "SHA-256",
            "MGF1",
            MGF1ParameterSpec.SHA256,
            PSource.PSpecified.DEFAULT
        )
    )
    val encryptedCekKey = Cipher.getInstance("RSA/ECB/OAEPWithSHA-256AndMGF1Padding").run {
        init(Cipher.ENCRYPT_MODE, publicKey, algParam)
        doFinal(cekKey.encoded)
    }

    val encodedEncryptedCekKey = encryptedCekKey.encodeUrlSafeBase64()
    val encodedIv = gcmSpec.iv.encodeUrlSafeBase64()
    val encodedCipherText = cipherTextBytes.encodeUrlSafeBase64()
    val encodedAuthTag = authTagBytes.encodeUrlSafeBase64()

    return "$header.$encodedEncryptedCekKey.$encodedIv.$encodedCipherText.$encodedAuthTag"
}

fun generateKeyPair(): KeyPair {
    val keyGen = KeyPairGenerator.getInstance("RSA")
    keyGen.initialize(2048)
    return keyGen.generateKeyPair()
}

private fun randomBytes(sizeInBits: Int) =
    ByteArray(sizeInBits / 8).apply {
        SecureRandom().nextBytes(this)
    }

private fun subArray(byteArray: ByteArray, beginIndex: Int, length: Int) =
    ByteArray(length).apply {
        System.arraycopy(byteArray, beginIndex, this, 0, size)
    }

fun ByteArray.encodeUrlSafeBase64() = String(Base64.getUrlEncoder().withoutPadding().encode(this))
fun String.decodeUrlSafeBase64() = Base64.getUrlDecoder().decode(this)

fun String.urlEncode() = URLEncoder.encode(this, "UTF-8")
fun String.urlDecode() = URLDecoder.decode(this, "UTF-8")

class Json constructor(init: Json.() -> Unit) : JSONObject() {
    init {
        init()
    }

    infix fun <T> String.to(value: T) {
        put(this, value)
    }
}