#!/usr/bin/ruby

=begin
Purpose:
  Processes GeoIP update requests (acts as local proxy to ensure
  we won't overload MaxMind's server). Requires refreshed database
  to be of any value.

Author: Wejn <wejn at box dot cz>
License: GPLv2 (without the "latter" option)
Requires: Ruby >= 1.8, geoip country license (to be of any value)
TS: 20060901184500
=end

# license keys we're gonna accept
$license_keys = {
	'lalala' => 'ns.wejn',
}

# db file we serve
$db_file = '/usr/share/GeoIP/GeoIP.dat'

require 'digest/md5'

module GeoIP
	class Update
		def initialize(file)
			@credentials = {}
			unless FileTest.exists?(file)
				raise ArgumentError, "db file doesn't exist"
			end
			@db_file = file
		end

		def credentials(enu)
			unless enu.kind_of?(Enumerable)
				raise ArgumentError, 'parameter not enumerable!'
			end

			enu.each do |key, *extra|
				@credentials[key] = true
			end

			true
		end

		def update_request(key, given_md5)
			raise 'License key invalid.' unless @credentials[key]

			content, md5 = nil, nil
			begin
				content = File.open(@db_file, 'r').read
				md5 = Digest::MD5.hexdigest(content)
			rescue
				raise 'WGIU/Error: Failed to load file and/or do checksum'
			end

			if given_md5 && given_md5.size >= 32 && md5 == given_md5
				raise 'No new updates available.'
			else
				content
			end
		end
	end
end

if __FILE__ == $0
	require 'cgi'
	require 'webrick'

	raise "DB file doesn't exist!" unless FileTest.exists?($db_file)

	gi = GeoIP::Update.new($db_file)
	gi.credentials($license_keys) if $license_keys.kind_of?(Enumerable)

	if ENV['GATEWAY_INTERFACE'] =~ /^cgi\/\d+.\d+$/i
		c = CGI.new
		out, parm = nil, nil
		begin
			out = gi.update_request(c.params['license_key'][0], c.params['md5'][0])
			parm = {
				'type' => 'application/octet-stream',
				'Content-Disposition' => 'attachment; filename=GeoIP.dat',
			}
		rescue
			out = $!.to_s
			parm = { 'type' => 'text/plain' }
		end
		c.out(parm) { out }
	elsif ARGV.first == '--server'
		port = ENV['LISTEN_ON'] =~ /^\d+$/ ? ENV['LISTEN_ON'].to_i : 
			(Process.uid.zero? ? 80 : 8000)
		s = WEBrick::HTTPServer.new(:Port => port)
		s.mount_proc('/app/update') do |req, res|
			begin
				res.body = gi.update_request(req.query['license_key'], req.query['md5'])
				res['Content-Type'] = 'application/octet-stream'
				res['Content-Disposition'] = 'attachment; filename=GeoIP.dat'
			rescue
				res.body = $!.to_s
				res['Content-Type'] = 'text/plain'
			end
		end
		trap("INT") { s.shutdown }
		s.start
	else
		$stderr.puts "Error: run either as CGI or give --server to launch webrick on 80(00)"
		exit 1
	end
end
