Add encrypted protocol related code, tests, samples (#4)

* Add encrypted protocol related code, tests, samples

* Fix ktlint

Author: Tougee

build.gradle

@@ -10,6 +10,7 @@ buildscript {
     ext.bouncycastleVersion = '1.69'
     ext.ed25519Version = '0.3.0'
     ext.curve25519Version = '0.5.0'
+    ext.ktlintVersion = '0.45.1'
     repositories {
         google()
         jcenter()
@@ -31,3 +32,8 @@ allprojects {
 task clean(type: Delete) {
     delete rootProject.buildDir
 }
+
+task ktlintFormat(group: "formatting") {
+    dependsOn ':library:ktlintFormat'
+    dependsOn ':samples:ktlintFormat'
+}

library/build.gradle

@@ -1,11 +1,10 @@
 plugins {
     id 'java-library'
     id 'kotlin'
+    id 'maven'
+    id 'java'
 }
 
-apply plugin: 'maven'
-apply plugin: 'java'
-
 java {
     sourceCompatibility = JavaVersion.VERSION_1_8
     targetCompatibility = JavaVersion.VERSION_1_8
@@ -29,7 +28,7 @@ dependencies {
     implementation "com.jakewharton.retrofit:retrofit2-kotlin-coroutines-adapter:${coroutineAdapterVersion}"
     implementation "net.i2p.crypto:eddsa:$ed25519Version"
     implementation "org.whispersystems:curve25519-java:$curve25519Version"
-    ktlint "com.pinterest:ktlint:0.40.0"
+    ktlint "com.pinterest:ktlint:$ktlintVersion"
 
     testImplementation 'org.jetbrains.kotlin:kotlin-test'
     testImplementation 'org.jetbrains.kotlin:kotlin-test-junit'

library/src/main/kotlin/one/mixin/bot/extension/UUIDExtension.kt

@@ -0,0 +1,11 @@
+package one.mixin.bot.extension
+
+import java.nio.ByteBuffer
+import java.util.UUID
+
+fun UUID.toByteArray(): ByteArray {
+    val bb = ByteBuffer.wrap(ByteArray(16))
+    bb.putLong(this.mostSignificantBits)
+    bb.putLong(this.leastSignificantBits)
+    return bb.array()
+}

library/src/main/kotlin/one/mixin/bot/util/ByteArrayUtil.kt

@@ -0,0 +1,14 @@
+package one.mixin.bot.util
+
+@ExperimentalUnsignedTypes
+fun toLeByteArray(v: UInt): ByteArray {
+    val b = ByteArray(2)
+    b[0] = v.toByte()
+    b[1] = (v shr 8).toByte()
+    return b
+}
+
+@ExperimentalUnsignedTypes
+fun leByteArrayToInt(bytes: ByteArray): UInt {
+    return bytes[0].toUInt() + (bytes[1].toUInt() shl 8)
+}

library/src/main/kotlin/one/mixin/bot/util/CryptoUtil.kt

@@ -1,8 +1,13 @@
+@file:Suppress("unused")
+
 package one.mixin.bot.util
 
 import net.i2p.crypto.eddsa.EdDSAPrivateKey
+import net.i2p.crypto.eddsa.EdDSAPublicKey
+import net.i2p.crypto.eddsa.math.FieldElement
 import net.i2p.crypto.eddsa.spec.EdDSANamedCurveTable
 import net.i2p.crypto.eddsa.spec.EdDSAPrivateKeySpec
+import net.i2p.crypto.eddsa.spec.EdDSAPublicKeySpec
 import one.mixin.bot.extension.base64Decode
 import one.mixin.bot.extension.base64Encode
 import org.bouncycastle.jce.provider.BouncyCastleProvider
@@ -36,9 +41,9 @@ fun generateEd25519KeyPair(): KeyPair {
     return net.i2p.crypto.eddsa.KeyPairGenerator().generateKeyPair()
 }
 
-@Throws(IllegalArgumentException::class)
-fun calculateAgreement(publicKey: ByteArray, privateKey: EdDSAPrivateKey): ByteArray =
-    Curve25519.getInstance(Curve25519.BEST).calculateAgreement(publicKey, privateKeyToCurve25519(privateKey.seed))
+fun calculateAgreement(publicKey: ByteArray, privateKey: ByteArray): ByteArray {
+    return Curve25519.getInstance(Curve25519.BEST).calculateAgreement(publicKey, privateKey)
+}
 
 fun privateKeyToCurve25519(edSeed: ByteArray): ByteArray {
     val md = MessageDigest.getInstance("SHA-512")
@@ -60,20 +65,39 @@ fun decryASEKey(src: String, privateKey: EdDSAPrivateKey): String? {
     return Base64.getEncoder().encodeToString(
         calculateAgreement(
             Base64.getUrlDecoder().decode(src),
-            privateKey
+            privateKeyToCurve25519(privateKey.seed)
         )
     )
 }
 
 private val secureRandom: SecureRandom = SecureRandom()
-private val GCM_IV_LENGTH = 12
+private const val GCM_IV_LENGTH = 12
 
 fun generateAesKey(): ByteArray {
     val key = ByteArray(16)
     secureRandom.nextBytes(key)
     return key
 }
 
+fun publicKeyToCurve25519(publicKey: EdDSAPublicKey): ByteArray {
+    val p = publicKey.abyte.map { it.toInt().toByte() }.toByteArray()
+    val public = EdDSAPublicKey(EdDSAPublicKeySpec(p, ed25519))
+    val groupElement = public.a
+    val x = edwardsToMontgomeryX(groupElement.y)
+    return x.toByteArray()
+}
+private fun edwardsToMontgomeryX(y: FieldElement): FieldElement {
+    val field = ed25519.curve.field
+    var oneMinusY = field.ONE
+    oneMinusY = oneMinusY.subtract(y)
+    oneMinusY = oneMinusY.invert()
+
+    var outX = field.ONE
+    outX = outX.add(y)
+
+    return oneMinusY.multiply(outX)
+}
+
 fun aesGcmEncrypt(plain: ByteArray, key: ByteArray): ByteArray {
     val iv = ByteArray(GCM_IV_LENGTH)
     secureRandom.nextBytes(iv)
@@ -136,7 +160,7 @@ private fun stripRsaPrivateKeyHeaders(privatePem: String): String {
     val lines = privatePem.split("\n".toRegex()).dropLastWhile { it.isEmpty() }.toTypedArray()
     lines.filter { line ->
         !line.contains("BEGIN RSA PRIVATE KEY") &&
-            !line.contains("END RSA PRIVATE KEY") && !line.trim { it <= ' ' }.isEmpty()
+            !line.contains("END RSA PRIVATE KEY") && line.trim { it <= ' ' }.isNotEmpty()
     }
         .forEach { line -> strippedKey.append(line.trim { it <= ' ' }) }
     return strippedKey.toString().trim { it <= ' ' }

library/src/main/kotlin/one/mixin/bot/util/EncryptedProtocol.kt

@@ -0,0 +1,73 @@
+package one.mixin.bot.util
+
+import net.i2p.crypto.eddsa.EdDSAPrivateKey
+import net.i2p.crypto.eddsa.EdDSAPublicKey
+import net.i2p.crypto.eddsa.spec.EdDSAPublicKeySpec
+import one.mixin.bot.extension.toByteArray
+import java.util.UUID
+
+class EncryptedProtocol {
+
+    @ExperimentalUnsignedTypes
+    fun encryptMessage(
+        privateKey: EdDSAPrivateKey,
+        plaintext: ByteArray,
+        otherPublicKey: ByteArray,
+        otherSessionId: String,
+        extensionSessionKey: ByteArray? = null,
+        extensionSessionId: String? = null
+    ): ByteArray {
+        val aesGcmKey = generateAesKey()
+        val encryptedMessageData = aesGcmEncrypt(plaintext, aesGcmKey)
+        val messageKey = encryptCipherMessageKey(privateKey.seed, otherPublicKey, aesGcmKey)
+        val messageKeyWithSession = UUID.fromString(otherSessionId).toByteArray().plus(messageKey)
+        val pub = EdDSAPublicKey(EdDSAPublicKeySpec(privateKey.a, ed25519))
+        val senderPublicKey = publicKeyToCurve25519(pub)
+        val version = byteArrayOf(0x01)
+
+        return if (extensionSessionKey != null && extensionSessionId != null) {
+            version.plus(toLeByteArray(2.toUInt())).plus(senderPublicKey).let {
+                val emergencyMessageKey =
+                    encryptCipherMessageKey(privateKey.seed, extensionSessionKey, aesGcmKey)
+                it.plus(UUID.fromString(extensionSessionId).toByteArray().plus(emergencyMessageKey))
+            }.plus(messageKeyWithSession).plus(encryptedMessageData)
+        } else {
+            version.plus(toLeByteArray(1.toUInt())).plus(senderPublicKey)
+                .plus(messageKeyWithSession)
+                .plus(encryptedMessageData)
+        }
+    }
+
+    @ExperimentalUnsignedTypes
+    fun decryptMessage(privateKey: EdDSAPrivateKey, sessionId: ByteArray, ciphertext: ByteArray): ByteArray {
+        val sessionSize = leByteArrayToInt(ciphertext.slice(IntRange(1, 2)).toByteArray()).toInt()
+        val senderPublicKey = ciphertext.slice(IntRange(3, 34)).toByteArray()
+        var key: ByteArray? = null
+        repeat(sessionSize) {
+            val offset = it * 64
+            val sid = ciphertext.slice(IntRange(35 + offset, 50 + offset)).toByteArray()
+            if (sessionId.contentEquals(sid)) {
+                key = ciphertext.slice(IntRange(51 + offset, 98 + offset)).toByteArray()
+            }
+        }
+        val messageKey = requireNotNull(key)
+        val message = ciphertext.slice(IntRange(35 + 64 * sessionSize, ciphertext.size - 1)).toByteArray()
+        val iv = messageKey.slice(IntRange(0, 15)).toByteArray()
+        val content = messageKey.slice(IntRange(16, messageKey.size - 1)).toByteArray()
+        val decodedMessageKey = decryptCipherMessageKey(privateKey.seed, senderPublicKey, iv, content)
+
+        return aesGcmDecrypt(message, decodedMessageKey)
+    }
+
+    private fun encryptCipherMessageKey(seed: ByteArray, publicKey: ByteArray, aesGcmKey: ByteArray): ByteArray {
+        val private = privateKeyToCurve25519(seed)
+        val sharedSecret = calculateAgreement(publicKey, private)
+        return aesEncrypt(sharedSecret, aesGcmKey)
+    }
+
+    private fun decryptCipherMessageKey(seed: ByteArray, publicKey: ByteArray, iv: ByteArray, ciphertext: ByteArray): ByteArray {
+        val private = privateKeyToCurve25519(seed)
+        val sharedSecret = calculateAgreement(publicKey, private)
+        return aesDecrypt(sharedSecret, iv, ciphertext)
+    }
+}