require 'webrick'
include WEBrick

require 'openid/consumer'
require 'openid/association'
require 'openid/interface'

module OpenID
	# First we need to define an AssociationManager to keep track of all
	# associations. This one is basically a copy of the one from httpconsumer.py
	# in the python module examples. It is not well designed, but works for an
	# example. I placed this in the OpenID namespace so it has easy access to
	# everything from the modules, without having to pull everything in.
	# Maybe that is bad form.
	class DictionaryAssociationManager < BaseAssociationManager
		def initialize()
			associator = DiffieHelmanAssociator.new(OpenID::SimpleHTTPClient.new)
			super(associator)
			@associations = []
		end
		
		def update(new_assoc, expired)
			@associations.push(new_assoc) if new_assoc
			
			if expired
				expired.each { |assoc1|
					@associations.each_index { |i|
						if assoc1 == @associations[i]
							@associations.delete_at(i)
							# break
						end
					}
				}
			end
		end
		
		def get_all(server_url)
			results = []
			@associations.each { |assoc|
				results.push(assoc) if assoc.server_url == server_url
			}
			return results
		end
		
		def invalidate(server_url, assoc_handle)
			@associations.each_index { |i|
				if @associations[i].server_url == server_url and @associations[i].handle == assoc_handle
					@associations.delete_at(i)
					# break
				end
			}
		end
	end
end #module OpenID
		
class OpenIDServlet < HTTPServlet::AbstractServlet

	def initialize(server, *options)
		@http_client = OpenID::SimpleHTTPClient.new
		@association_manager = OpenID::DictionaryAssociationManager.new
		@consumer = OpenID::Consumer.new(@http_client, @association_manager)
		super(server, *options)
	end
	
	def do_GET(req, res)
		res['Content-Type'] = 'text/html'
		if req.query['identity_url']
			# This executes after the form has been submitted.
			identity_url = req.query['identity_url'].to_s
			redirect_url = @consumer.handle_request(identity_url, req.header['referer'][0].to_s)
			if redirect_url
				# Redirect to the authentication page (The openid server)
				res.set_redirect(HTTPStatus::Found, redirect_url)
			else
				res.status = 412
				res.body = "no identity url!"
			end
		elsif req.query['openid.mode']
			# This is where we handle the user after they are sent back by the server
			begin
				valid_to = @consumer.handle_response(OpenID::Request.new(req.query, 'GET'))
			rescue
				# You can rescue specific errors here to deal with the user cancelling and such.
				res.body = get_output_page("Exception raised by consumer.handle_response:\n<br />" + $!)
			else
				res.body = get_output_page("Logged in! Valid-to: #{valid_to}")
			end
		else
			# Display the login form
			res.body = get_input_form
		end
	end
	
	def get_input_form()
		return <<EOF
		<html>
		<head><title>OpenID for Ruby</title></head>
		<body style="background-color: #FFFFCC;">
		<h2>OpenID for Ruby httpconsumer.rb example</h2>
		<form method="GET" action="/">
		Your Identity URL: <input type="text" name="identity_url" size="60" />
		<br /><input type="submit" value="Log in" />
		</form>
		</body>
		</html>
EOF
	end
	
	def get_output_page(text, bgcolor = '#FFFFCC')
		return <<EOF
		<html>
		<head><title>OpenID for Ruby</title></head>
		<body style="background-color: #{bgcolor};">
		<h2>OpenID for Ruby httpconsumer.rb example</h2>
		#{text}
		</body>
		</html>
EOF
	end

end

server = HTTPServer.new( :Port => 8081 )
server.mount('/', OpenIDServlet)
trap("INT") { server.shutdown }
server.start