#!/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
                raise 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