Skip to content

Instantly share code, notes, and snippets.

Last active May 7, 2020 23:50
Show Gist options
  • Save Brunomachadob/34f39128df8fef369a647bdbb8f6d44b to your computer and use it in GitHub Desktop.
Save Brunomachadob/34f39128df8fef369a647bdbb8f6d44b to your computer and use it in GitHub Desktop.
Educational-purposed implementation of Diffie Hellman key exchange

This is a educational purposed implementation of the Diffie Hellman key exchange for encrypted conversation between two parties.

Do not use this in production (as it is simplifying a lot of stuff, despite I know you will ignore this)

If you use IntelliJ, you can just paste the implementation into a new scratch and run it.

import java.math.BigInteger
import javax.crypto.Cipher
import javax.crypto.spec.SecretKeySpec
import kotlin.system.measureTimeMillis
val random = SecureRandom()
// For brevity, public G and N are being pre-computed
val g: BigInteger = BigInteger.probablePrime(10, random)
val n: BigInteger = BigInteger.probablePrime(1000, random)
enum class MessageType {
class Message(val type: MessageType, val payload: ByteArray)
Using a simple AES algorithm with a symmetric key derived from the session-key
using a SHA-512 hash.
Apparently the usage of truncate(SHA-512, 32) is better than SHA-256
class AESCipher(keyBytes: ByteArray) {
companion object {
private const val ALGORITHM = "AES"
private val SHA512: MessageDigest = MessageDigest.getInstance("SHA-512")
private val encryptor: Cipher
private val decryptor: Cipher
init {
val keySpec = SHA512.digest(keyBytes)
.let { SecretKeySpec(it, ALGORITHM) }
encryptor = Cipher.getInstance(ALGORITHM).also { it.init(Cipher.ENCRYPT_MODE, keySpec) }
decryptor = Cipher.getInstance(ALGORITHM).also { it.init(Cipher.DECRYPT_MODE, keySpec) }
fun encrypt(message: ByteArray): ByteArray = encryptor.doFinal(message)
fun decrypt(message: ByteArray): ByteArray = decryptor.doFinal(message)
This class would represent a channel out in the public, for example the internet
where you can have people intercepting the messages going through it
inner class PublicChannel {
private val dnsServer = mutableMapOf<String, Actor>()
fun register(name: String, actor: Actor) = apply {
dnsServer[name] = actor
fun sendMessage(from: String, to: String, message: Message): Message {
val dest: Actor = dnsServer[to]
?: throw IllegalArgumentException("Actor $to does not exist")
println("PublicChannel intercepted a message: ${interceptedMessageAsString(message)}")
val response = dest.receiveMessage(from, message)
println("PublicChannel intercepted a response: ${interceptedMessageAsString(response)}")
return response
private fun interceptedMessageAsString(message: Message): Any =
"{type=${message.type}, message=" + when (message.type) {
MessageType.HANDSHAKE_HELLO -> BigInteger(message.payload)
MessageType.HANDSHAKE_FINISH -> String(message.payload)
MessageType.HANDSHAKE_ERROR -> "Handshake error"
MessageType.ENCRYPTED_MESSAGE -> String(message.payload)
} + "}"
This class represents two entities that are going to publicly exchange their public keys
and talk between each other using encrypted messages after the exchange
inner class Actor(
private val name: String,
private val channel: PublicChannel
) {
private val privateKey = BigInteger.probablePrime(n.bitLength() / 2, random)
private val publicKey = g.modPow(privateKey, n).toByteArray()
private var cipher: AESCipher? = null
init {
channel.register(name, this)
fun receiveMessage(from: String, message: Message): Message = when (message.type) {
MessageType.HANDSHAKE_HELLO -> handleHandshakeHello(message)
MessageType.HANDSHAKE_FINISH -> handleHandshakeFinish(from, message)
MessageType.HANDSHAKE_ERROR -> handleHandshakeError()
MessageType.ENCRYPTED_MESSAGE -> handleEncryptedMessage(message)
private fun createCipher(otherPublicKey: ByteArray) {
val sessionKey = BigInteger(otherPublicKey).modPow(privateKey, n)
cipher = AESCipher(sessionKey.toByteArray())
private fun handshake(to: String) {
// We send our public key and get the other public key
val othersPublicKey = channel
.sendMessage(name, to, Message(MessageType.HANDSHAKE_HELLO, publicKey)).payload
// As a finish message to verify the encryption, I'm sending simply "$name:finish"
// and validating it on the other side.
val response = channel
.sendMessage(name, to, Message(MessageType.HANDSHAKE_FINISH, buildHandshakeFinishPayload(name)))
if (response.type == MessageType.HANDSHAKE_FINISH) {
val expectedMessage = buildHandshakeFinishPayload(to)
if (!response.payload.contentEquals(expectedMessage)) {
throw IllegalStateException("Failed to verify HANDSHAKE_FINISH message")
} else throw IllegalStateException("Invalid during handshake state: ${response.type}")
private fun handleHandshakeHello(message: Message): Message {
return Message(MessageType.HANDSHAKE_HELLO, publicKey)
private fun buildHandshakeFinishPayload(name: String) =
"$name:finished".let { cipher!!.encrypt(it.toByteArray()) }
private fun handleHandshakeFinish(from: String, message: Message): Message {
val expectedMessage = buildHandshakeFinishPayload(from)
if (!expectedMessage.contentEquals(message.payload)) {
return Message(MessageType.HANDSHAKE_ERROR, byteArrayOf())
return Message(MessageType.HANDSHAKE_FINISH, buildHandshakeFinishPayload(name))
private fun handleHandshakeError(): Nothing = throw IllegalStateException("$name got a handshake error")
private fun handleEncryptedMessage(message: Message): Message {
val plainTextMessage = String(cipher!!.decrypt(message.payload))
println("$name got: $plainTextMessage")
return Message(message.type, cipher!!.encrypt("Hey there!".toByteArray()))
fun sendMessage(to: String, message: String): String {
if (cipher == null) handshake(to)
val encStr = cipher!!.encrypt(message.toByteArray())
val response = channel
.sendMessage(name, to, Message(MessageType.ENCRYPTED_MESSAGE, encStr))
return String(cipher!!.decrypt(response.payload))
val channel = PublicChannel()
val alice = Actor("alice", channel)
val bob = Actor("bob", channel)
alice.sendMessage("bob", "Hey Bob!")
bob.sendMessage("alice", "Hey Alice!")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment