DiffieHellmanMerkle.scala

object DiffieHellmanMerkle {

  import java.math.BigInteger
  import scala.language.implicitConversions
  import scala.util.Random

  def main(args: Array[String]): Unit =
    println(
      diffieHellmanMerkle(generator = 3, modulus = 17, alicePrivateKey = 54, bobPrivateKey = 24)
    )

  implicit def bigIntToInt(i: BigInteger): Int = i.intValue
  
  implicit def intToBigInt(i: Int): BigInteger = new BigInteger(i.toString)

  def probRandomPrime(maxBaseForNextPrime: Int = 100, rng: Random = new Random()): Int =
    math.abs(rng.nextInt(maxBaseForNextPrime)).nextProbablePrime.intValue

  case class DhResult(
                       generator: Int,
                       modulus: Int,
                       alicePrivateKey: Int,
                       bobPrivateKey: Int,
                       alicePublicKey: Int,
                       bobPublicKey: Int,
                       aliceSharedSecret: Int,
                       bobSharedSecret: Int
                       ) {

    override def toString: String =
      s"""DhResult(
          generator=$generator, modulus=$modulus, alicePrivateKey=$alicePrivateKey,
          bobPrivateKey=$bobPrivateKey, alicePublicKey=$alicePublicKey,
          bobPublicKey=$bobPublicKey, aliceSharedSecret=$aliceSharedSecret,
          bobSharedSecret=$bobSharedSecret
          )"""
  }

  def diffieHellmanMerkle(
                     generator: Int = probRandomPrime(maxBaseForNextPrime = 10),
                     modulus: Int = probRandomPrime(maxBaseForNextPrime = 20),
                     alicePrivateKey: Int,
                     bobPrivateKey: Int
                     ): DhResult = {

    val alicePublicKey = generator.pow(alicePrivateKey).mod(modulus)
    val bobPublicKey = generator.pow(bobPrivateKey).mod(modulus)
    val aliceSharedSecret = bobPublicKey.pow(alicePrivateKey).mod(modulus)
    val bobSharedSecret = alicePublicKey.pow(bobPrivateKey).mod(modulus)
    DhResult(
      generator, modulus, alicePrivateKey, bobPrivateKey,
      alicePublicKey, bobPublicKey,
      aliceSharedSecret, bobSharedSecret
    )
  }
}