Diffie-Hellman in Ruby

I looked around and couldn’t find a pure-ruby implementation of Diffie-Hellman key exchange. Diffie-Hellman key exchange is a nifty way to end up with the same shared secret between Alice and Bob without ever sending the secret key to the other side. It’s used in ISAKMP, SSH and a host of other crypto-based protocols. The code for Diffie-Hellman in Ruby is unbelievably terse to the point you wonder if you actually got it working right. Two things come in handy: Ruby has open classes that you can extend and Ruby has built-in Bignum support. Integers don’t overflow in Ruby, they just keep expanding.

We are going to need to extend the built-in Integer class with two methods. The first one does modular-exponentiation, while the second one counts the number of bits set in a given number. Since Fixnum’s and Bignum’s in Ruby derive from Integer, these two methods are automatically available for all numbers.

First we extend the Integer class:

class Integer
    # Compute self ^ e mod m
    def mod_exp e, m
        result = 1
        b = self
        while e > 0
            result = (result * b) % m if e[0] == 1
            e = e >> 1
            b = (b * b) % m
        end
        return result
    end

    # A roundabout, slow but fun way of counting bits.
    def bits_set
        ("%b" % self).count('1')
    end
end

and here’s the Diffie-Hellman class:

class DH
    attr_reader :p, :g, :q, :x, :e

    # p is the prime, g the generator and q order of the subgroup
    def initialize p, g, q
        @p = p
        @g = g
        @q = q
    end

    # generate the [secret] random value and the public key
    def generate tries=16
        tries.times do
            @x = rand(@q)
            @e = self.g.mod_exp(@x, self.p)
            return @e if self.valid?
        end
        raise ArgumentError, "can't generate valid e"
    end

    # validate a public key
    def valid? _e = self.e
        _e and _e.between?(2, self.p-2) and _e.bits_set > 1
    end

    # compute the shared secret, given the public key
    def secret f
        f.mod_exp(self.x, self.p)
    end
end

and finally, here’s how you use it

alice = DH.new(53, 5, 23)
bob   = DH.new(53, 5, 15)
alice.generate
bob.generate

alice_s = alice.secret(bob.e)
bob_s   = bob.secret(alice.e)
puts alice_s
puts bob_s

alice_s and bob_s should/will be the same.