GeoIP update client + server

#!/usr/bin/ruby

=begin
Purpose:
  Updates GeoIP Country (paid) database from MaxMind.com. Requires
  valid license key to work.

Why:
  1) I hate the idea of running some weird C code ("geoipupdate")
     from cron, that's why.
  2) I don't want to run "geoipupdate" on all machines using this DB.

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: 20060901171500
Changelog:
  20060626181500 - Initial release.
  20060901081500 - Invalid license key throwed as exception.
                   Update url (base) configurable.
  20060901171500 - Added custom error check (compatibility w/ my geoup server)
=end

# your license key here
$license_key = 'You-wish----'

# destination file
$outfile = '/usr/share/GeoIP/GeoIP.dat'

# temp file -- MUST BE on same filesystem as $outfile
$tmpfile = '/usr/share/GeoIP/GeoIP.dat.tmp'

# base url
$baseurl = 'ht' + 'tp://www.maxmind.com/app/update'

require 'digest/md5'
require 'open-uri'
require 'zlib'
require 'stringio'

# Fetch GeoIP country file from MaxMind
def get_result_for(key, md5)
  url = $baseurl + "?license_key=#{key}&md5=#{md5}"

  content = nil
  open(url) do |io|
    compr = StringIO.new(raw = io.read)
    if compr.read(2) == "\x1f\x8b"
      compr.rewind
      gz = Zlib::GzipReader.new(compr)
      content = gz.read
      gz.close
    else
      content = raw
    end
  end

  raise "invalid license key" if content =~ /License key invalid/i
  raise "Remote: #{content.sub(/.*? /, '')}" if content =~ /^WGIU\/Error: /i

  content
end

puts "Begin." if $VERBOSE

begin
  omd5sum = Digest::MD5.hexdigest(File.open($outfile, 'r').read) rescue nil
  content = get_result_for($license_key, omd5sum)

  if content =~ /No new updates available/i
    puts "No new updates ..." if $VERBOSE
    exit 0
  end

  md5sum = Digest::MD5.hexdigest(content)
  content2 = get_result_for($license_key, md5sum)
rescue
  puts "Error:  Problem fetching file! >> #{$!} (#{$!.class})"
  exit 1
end

puts "Fetched." if $VERBOSE

if content2 =~ /No new updates available/
  begin
    File.open($tmpfile, 'w') do |f|
      f.write(content)
    end
    File.unlink($outfile) if FileTest.exists?($outfile)
    File.rename($tmpfile, $outfile)
  rescue
    puts "Error: Problem updating file! >> #{$!} (#{$!.class})"
    exit 1
  end
else
  puts "Error: Problem refreshing GeoIP -- new file fails MD5 check!"
  exit 1
end

puts "Updated." if $VERBOSE
exit 0

#!/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"
    exit 1
  end
end