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')
        #to_s(2).count('1')   # alternative
        #count = 0         # alternative
        #byte = self.abs
        #count += byte & 1 and byte >>= 1 until byte == 0     # cf. http://snippets.dzone.com/posts/show/4233
        #count
    end
end


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

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

__author__="Andrew Pennebaker (andrew.pennebaker@gmail.com)"
__date__="21 Dec 2005 - 17 Jul 2006"
__copyright__="Copyright 2006 Andrew Pennebaker"
__license__="GPL"
__version__="0.3"
__URL__="http://snippets.dzone.com/posts/show/2888"

import sys
from getopt import getopt

TEST_MODE="TEST"
HASH_MODE="HASH"

class HashFunction:
	"""Base class for hash functions"""

	# bytes
	BLOCK_SIZE=1
	DIGEST_SIZE=1

	INIT=0x00
	SUM_REQ="Sum >= 0"

	TEST_DATA="abc"
	TEST_HASH=0x26

	def __init__(self, sum=0x00):
		"""Sets the context to some initial sum"""
		self.sum=sum&0xff

	def sumValid(self, sum):
		return sum>=0

	def _update(self, b):
		"""Data is an array of bytes"""

		self.sum=(self.sum+b)&0xff

	def update(self, data):
		"""Helper for _update(). Data is a string."""

		for byte in data:
			self._update(ord(byte))

	def digest(self):
		"""Returns an integer or long integer"""

		return self.sum

	def formatDigest(self):
		"""Returns a formatted string, different for each algorithm"""

		return "%02x" % (self.digest())

	def unformatDigest(self, hash):
		"""Converts formatted hash into integer"""

		return int(hash, 16)

	def reset(self):
		"""Reinitializes the hash"""

		self.__init__()

	def test(self):
		self.reset()
		self.update(self.TEST_DATA)
		hash=self.digest()
		formattedHash=self.formatDigest()
		unformattedHash=self.unformatDigest(formattedHash)

		if unformattedHash==self.TEST_HASH:
			return "OK"

		return [self.TEST_DATA, hash, formattedHash, unformattedHash, self.TEST_HASH]

def usage(script=sys.argv[0]):
	print "Usage: %s [options] <file1 file2 file3 ... >" % (script)
	print "\n-s --sum <sum>"
	print "-t --test engine"
	print "-h --help usage"

	sys.exit()

def main(hasher=HashFunction, script="HashFunction.py"):
	global TEST_MODE
	global HASH_MODE

	systemArgs=sys.argv[1:] # ignore program name

	mode=HASH_MODE
	sum=hasher.INIT

	optlist, args=[], []
	try:
		optlist, args=getopt(systemArgs, "s:th", ["sum=", "test", "help"])
	except Exception, e:
		usage(script)

	if len(optlist)<1 and len(args)<1:
		usage(script)

	for option, value in optlist:
		if option=="-h" or option=="--help":
			usage(script)
		elif option=="-t" or option=="--test":
			mode=TEST_MODE
		elif option=="-s" or option=="--sum":
			try:
				sum=hasher.unformatDigest(value)
				if not hasher.sumValid(sum):
					raise Exception
			except Exception, e:
				raise "Requires: %s" % (hasher.SUM_REQ)

	hasher=hasher()

	if mode==TEST_MODE:
		print hasher.test()
	elif mode==HASH_MODE:
		for file in args:
			f=None
			try:
				f=open(file, "rb")
			except Exception, e:
				print "Could not open %s" % (file)
			else:
				hasher.reset()
				for line in f:
					hasher.update(line)
				f.close()

				print hasher.formatDigest()

if __name__=="__main__":
	main(HashFunction, "HashFunction.py")