require 'openid/constants'
require 'openid/errors'
require 'openid/util'

require 'base64'
require 'digest/sha1'

module OpenID
	class Association
		attr_reader :handle, :secret, :expiry
		attr_writer :handle, :secret, :expiry
		def initialize(handle, secret, expiry, replace_after)
			@handle = handle.to_s
			@secret = secret.to_s
			@replace_after = replace_after
			@expiry = expiry
		end
		#TODO: override ==?
	end #class Association
	
	# Represents an association established for a consumer.
	class ConsumerAssociation < Association
		attr_reader :server_url
		attr_writer :server_url
		def initialize(server_url, handle, secret, expiry, replace_after)
			super(handle, secret, expiry, replace_after)
			@server_url = server_url.to_s
		end
		
		def get_replace_after
			if @replace_after
				return @replace_after
			else
				return @expiry
			end
		end
	end #class ConsumerAssociation
	
	class ServerAssociation < Association
		def initialize(handle, secret, expiry_off, replace_after_off)
			now = Time.now
			expiry = now + expiry_off
			replace_after = now + replace_after_off
			super(handle, secret, expiry, replace_after)
			@issued = now
		end
	end #class ServerAssociation

	# Abstract class which is the parent of both DumbAssociationManager and 
	# BaseAssociationManager.
	class AssociationManager
		# Return a ConsumerAssociation based on server_url and assoc_handle
		def get_association(server_url, assoc_handle)
			raise NotImplementedError
		end
		# Create an association with server_url. Return an assoc_handle
		def associate(server_url)
			raise NotImplementedError
		end
		# Invalidate an Association
		def invalidate(server_url, assoc_handle)
			raise NotImplementedError
		end
	end # class AssociactionManager
		
	class DumbAssociationManager < AssociationManager
		def get_association(server_url, assoc_handle)
			return nil
		end
		
		def associate(server_url)
			return nil
		end
		
		def invalidate(server_url, assoc_handle)
			# pass
		end
	end #class DumbAssociationManager
		
	class BaseAssociationManager < AssociationManager
		def initialize(associator)
			@associator = associator
		end
		# Returns assoc_handle associated with server_url
		def associate(server_url)
			now = Time.now
			expired = []
			assoc = nil
			get_all(server_url).each { |current|
				replace_after = current.get_replace_after
				if current.expiry < now
					expired.push(current)
				elsif assoc == nil
					assoc = current if replace_after > now
				elsif replace_after > assoc.replace_after
					assoc = current
				end
			}
			new_assoc = nil
			if assoc == nil
				assoc = new_assoc = @associator.associate(server_url)
			end
			if new_assoc or expired
				update(new_assoc, expired)
			end
			
			return assoc.handle
		end
		
		def get_association(server_url, assoc_handle)
			get_all(server_url).each { |assoc|
				return assoc if assoc.handle == assoc_handle
			}
			return nil
		end
		
		# This must be implemented by subclasses.
		#
		# new_assoc is either a new association object or nil. Expired is a possibly
		# empty list of expired associations.
		# Subclass should add new_assoc if it is not nil, and expire each association
		# in the expired list.
		def update(new_assoc, expired)
			raise NotImplementedError
		end
		
		# This must be implemented by subclasses.
		#
		# Should return a list of Association objects matching server_url
		def get_all(server_url)
			raise NotImplementedError
		end
		
		# This must be implemented by subclasses.
		#
		# Subclass should remove the association for the given
		# server_url and assoc_handle.
		def invalidate(server_url, assoc_handle)
			raise NotImplementedError
		end
	end #class BaseAssociationManager	
	
	# A class for establishing associations with OpenID servers.
	class DiffieHelmanAssociator
		include OpenID
		def initialize(http_client)
			@http_client = http_client
		end
		
		# Returns modulus and generator for Diffie-Helman.
		# Override this for non-default values.
		def get_mod_gen()
			return DEFAULT_DH_MODULUS, DEFAULT_DH_GEN
		end
		
		# Establishes an association with an OpenID server, indicated by server_url.
		# Returns a ConsumerAssociation.
		def associate(server_url)
			p, g = get_mod_gen()
			private_key = (1 + rand(p-2))
			dh_public = g.mod_exp(private_key, p)
			args = {
				'openid.mode' => 'associate',
				'openid.assoc_type' => 'HMAC-SHA1',
				'openid.session_type' => 'DH-SHA1',
				'openid.dh_modulus' => Base64.encode64(p.to_btwoc).delete("\n"),
				'openid.dh_gen' => Base64.encode64(g.to_btwoc).delete("\n"),
				'openid.dh_consumer_public' => Base64.encode64(dh_public.to_btwoc).delete("\n")
			}
			body = url_encode(args)
			url, data = @http_client.post(server_url, body)
			now = Time.now.utc
			# results are temporarily stored in an instance variable.
			# I'd like to restructure to avoid this.
			@results = parse_kv(data)
			
			assoc_type = get_result('assoc_type')
			if assoc_type != 'HMAC-SHA1'
				raise RuntimeError, "Unknown association type: #{assoc_type}", caller
			end
			
			assoc_handle = get_result('assoc_handle')
			issued = DateTime.strptime(get_result('issued')).to_time
			expiry = DateTime.strptime(get_result('expiry')).to_time
			
			delta = now - issued
			expiry = expiry + delta
			
			replace_after_s = @results['replace_after']
			if replace_after_s
				replace_after = DateTime.strptime(replace_after_s).to_time + delta
			else
				replace_after = nil
			end
			
			session_type = @results['session_type']
			if session_type
				if session_type != 'DH-SHA1'
					raise RuntimeError, "Unknown Session Type: #{session_type}", caller
				end
				dh_server_pub = from_btwoc(Base64.decode64(get_result('dh_server_public')))
				enc_mac_key = get_result('enc_mac_key')
				dh_shared = dh_server_pub.mod_exp(private_key, p)

				secret = Base64.decode64(enc_mac_key) ^ Digest::SHA1.digest(dh_shared.to_btwoc)
			else
				secret = get_result('mac_key')
			end
			@results = nil
			return ConsumerAssociation.new(server_url, assoc_handle, secret, expiry, replace_after)
		end	
	private
		def get_result(key)
			begin
				return @results.fetch(key)
			rescue IndexError
				raise ProtocolError, "Association server response missing argument #{key}", caller
			end
		end
	end #class DiffieHelmanAssociator
		
end #module OpenID