#!ruby1.8

require "rubygems"
require "mongrel"
require "mongrel/handlers"
require "webrick/httputils"
require "net/http"
require "uri"
require 'yaml'
require "pp"
require "pathname"
require "stringio"
require "zlib"
require "optparse"
require "tmpdir"
require "digest/md5"

class Pathname
	@@tempname_number = 0
	def self.tempname(base=$0, dir=Dir.tmpdir)
		@@tempname_number += 1
		name = "#{dir}/#{File.basename(base)}.#{$$}.#{@@tempname_number}"
		path = new(name)
		at_exit do
			path.rmtree if path.exist?
		end
		path
	end
end



class CocProxyCommand
	VERSION = "$Revision$"
	DEFAULT_CONFIG = {
		:Port        => 5432,
		:FilterDir   => "files",
		:Rules       => [
			"\#{File.basename(req.path_info)}",
			"\#{req.host}\#{req.path_info}",
			"\#{req.host}/\#{File.basename(req.path_info)}",
			".\#{req.path_info}",
		]
	}

	def self.run(argv)
		new(argv.dup).run
	end

	def initialize(argv)
		@argv = argv
		@parser = OptionParser.new do |parser|
			parser.banner = <<-EOB.gsub(/^\t+/, "")
				Usage: #$0 [options]
			EOB

			parser.separator ""
			parser.separator "Options:"

			parser.on("-c", "--config CONFIG.yaml") do |config|
				begin
					@config = YAML.load_file(config)
				rescue Errno::ENOENT
					puts "#{config} is not found"
					exit
				end
			end

			parser.on("-p", "--port PORT", "Specify port number. This option overrides any config.") do |port|
				@port = port.to_i
			end

			parser.on("-n", "--no-cache", "Disable cache.") do |port|
				@nocache = true
			end

#			parser.on("--disable-double-screen", "Disable loading double_screen.rb") do |c|
#				@disable_double_screen = c
#			end

			parser.on("--version", "Show version string `#{VERSION}'") do
				puts VERSION
				exit
			end
		end
	end

	def run
		@parser.order!(@argv)
		$stdout.sync = true

		unless @config
			begin
				@config = YAML.load_file("proxy-config.yaml")
				puts "proxy-config.yaml was found. Use it."
			rescue Errno::ENOENT
				@config = {
					"server" => {
					},
				}
				puts "Use default configuration."
			end
		end

		server_config = DEFAULT_CONFIG.update(@config["server"])
		server_config[:Port]    = @port if @port
		server_config[:nocache] = @nocache

#		unless @disable_double_screen
#			begin
#				require "double_screen.rb"
#			rescue LoadError => e
#			end
#		end
		puts "Port : #{server_config[:Port]}"
		puts "Dir  : #{server_config[:FilterDir]}/"
		puts "Cache: #{!server_config[:nocache]}"
		puts "Rules:"
		server_config[:Rules].each_with_index do |item, index|
			puts "    #{index+1}. #{item}"
		end

		server = Mongrel::HttpServer.new('0.0.0.0', server_config[:Port])
		server.register("/", ArrogationProxyServer.new(server_config))
		server.run
		sleep
	end

	class ProxyHandler < Mongrel::HttpHandler
		include Mongrel::Const

		def process(req, res)
			uri = URI(req.params["REQUEST_URI"])
			header = choose_header(Hash[*req.params.map {|k, v|
				k = k.sub(/^HTTP_/, "")
				Regexp.last_match ? [k.tr("_", "-"), v] : nil
			}.compact.flatten])
			begin
				response = nil
				Net::HTTP.new(uri.host, uri.port).start do |http|
					http.open_timeout = 3   # secs  #   necessary (maybe bacause
					http.read_timeout = 6   # secs  #   Ruby's bug, but why?)
					response = http.send_request(req.params["REQUEST_METHOD"], uri.request_uri, req.body.read, header)
				end
				filter(req, response)
				res.start(response.code.to_i, false, response.message) do |head, out|
					choose_header response, head
					out << response.body
				end
			rescue => e
				res.start(500, false, "Exception raised") do |head, out|
					head["Content-Type"] = "text/plain"
					out << e.message << "\n"
					out << e.backtrace.join("\n\t")
				end
			end
		end

		# from WEBrick
		HopByHop = %w( CONNECTION KEEP-ALIVE PROXY-AUTHENTICATE UPGRADE PROXY-AUTHORIZATION TE TRAILERS TRANSFER-ENCODING )
		ShouldNotTransfer = %w( PROXY-CONNECTION )

		def choose_header(src, dst={})
			connections = (src['connection'] || "").split(/,\s+/).collect{|i| i.downcase }
			src.each { |key, value|
				key = key.upcase
				if HopByHop.member?(key)          || # RFC2616: 13.5.1
				   connections.member?(key)       || # RFC2616: 14.10
				   ShouldNotTransfer.member?(key)    # pragmatics

					next
				end
				dst[key] = value
			}
			dst
		end

		def filter(req, res)
		end
	end

	class ArrogationProxyServer < ProxyHandler
		def initialize(config)
			@config = config
			@cache_dir = Pathname.tempname
			@cache_dir.mkpath
		end

		def process(req, res)
			dir = @config[:FilterDir]
			def req.path_info
				r = params["PATH_INFO"]
				r == "/" ? "" : r
			end
			def req.host
				params["HTTP_HOST"]
			end
			def req.request_uri
				@request_uri ||= URI(params["REQUEST_URI"])
			end
			def req.query
				request_uri.query || ""
			end
			$stderr.puts req.path_info if $DEBUG

			content    = ""
			local_path = ""

			@config[:Rules].each do |path|
				path = "#{dir}/#{eval("%Q(#{path})")}"
				# $stderr.puts "Checking #{path.to_s}"
				if FileTest.file? path
					puts "Hit Arrogation: #{req.path_info}"
					local_path = path
					content = File.open(path).binmode.read
					break
				end
			end

			req.params.delete("HTTP_IF_MODIFIED_SINCE")

			case
			when content =~ /proxy-replace:\s*(.+)\s*/
				@content = content
				puts "#{local_path} =>"
				super
			when content !~ /\A\s*\Z/
				mime_types = WEBrick::HTTPUtils::DefaultMimeTypes.update(@config[:MimeTypes] || {})
				res.start do |head, body|
					head["Content-Type"] = WEBrick::HTTPUtils.mime_type(req.path_info, mime_types)
					body << content
				end
				puts "Rewrote: <= #{local_path}"
			else
				if @config[:nocache]
					super
				else
					path = @cache_dir + Digest::MD5.hexdigest(req.params["REQUEST_URI"])
					if path.exist?
						puts "Hit Cache: #{req.params["REQUEST_URI"]}"
						response = path.open("rb") {|i| Marshal.load(i) }
						res.start(response.code.to_i, false, response.message) do |head, out|
							choose_header response, head
							out << response.body
						end
					else
						super
					end
				end
			end
		rescue Exception => e
			puts e.message
			puts e.backtrace
			puts
		end

		def filter(req, res)
			if @content
				@content.sub!(/proxy-replace:\s*(.+)\s*/, "")
				regexp = Regexp.new(Regexp.last_match[1])
				puts "Replace Regexp: #{regexp.source}"

				case (res["Content-Encoding"] || "").downcase
				#when "deflate"
				#when "compress"
				when "gzip"
					res["Content-Encoding"] = nil
					res.body.replace Zlib::GzipReader.wrap(StringIO.new(res.body)) {|gz| gz.read }
				end

				p res

				m = res.body.match(regexp)
				if m && m[1]
					res.body[m.begin(1)..(m.end(1)-1)] = @content
				else
					puts "In-place Regexp match failed..."
				end
				res["Content-Length"] = res.body.length.to_s
				@content = nil
			end

			unless @config[:nocache]
				path = @cache_dir +  Digest::MD5.hexdigest(req.params["REQUEST_URI"])
				path.open("wb") do |f|
					Marshal.dump(res, f)
				end
			end
		end
	end

end

CocProxyCommand.run(ARGV)