#!/usr/bin/ruby

require "apache_log_regex"
require "date"
require "file-tail"
require "gdbm"
require "lru_redux"

REQUESTS_PER_SECOND = <%= node[:tile][:ratelimit][:requests_per_second] %>
BLOCK_AT = <%= node[:tile][:ratelimit][:maximum_backlog] %>
UNBLOCK_AT = BLOCK_AT / 2

parser = ApacheLogRegex.new('%a %l %u %t \"%r\" %>s %O \"%{Referer}i\" \"%{User-Agent}i\"')
clients = LruRedux::Cache.new(1000000)

def decay_count(client, time)
  decay = (time.to_i - client[:last_update]) * REQUESTS_PER_SECOND

  client[:request_count] = [client[:request_count] - decay, 0].max
  client[:last_update] = time.to_i
end

def write_blocked_ips(clients)
  time = Time.now

  File.open("/srv/tile.openstreetmap.org/conf/ip.map.new", "w") do |file|
    clients.each do |address, client|
      decay_count(client, time)

      if client[:request_count] >= UNBLOCK_AT
        file.puts "#{address} blocked"
      elsif client.has_key?(:blocked_at)
        puts "Unblocked #{address}"

        client.delete(:blocked_at)
      end
    end
  end

  File.rename("/srv/tile.openstreetmap.org/conf/ip.map.new",
              "/srv/tile.openstreetmap.org/conf/ip.map")

  time + 900
end

next_check = write_blocked_ips(clients)

File::Tail::Logfile.tail("/var/log/apache2/access.log") do |line|
  begin
    hash = parser.parse!(line)

    address = hash["%a"]
    request = hash["%r"]

    next if address == "127.0.0.1" || address == "::1"

    time = Time.now

    client = clients.getset(address) do
      { :request_count => 0, :last_update => 0 }
    end

    decay_count(client, time)

    if request =~ %r{^(GET|POST) /cgi-bin/export.*}
      client[:request_count] = client[:request_count] + 150
    else
      client[:request_count] = client[:request_count] + 1
    end

    if client[:request_count] > BLOCK_AT && !client.has_key?(:blocked_at)
      puts "Blocked #{address}"

      client[:blocked_at] = time

      next_check = time
    elsif client[:request_count] < UNBLOCK_AT && client.has_key?(:blocked_at)
      puts "Unblocked #{address}"

      client.delete(:blocked_at)

      next_check = time
    end

    if time >= next_check
      next_check = write_blocked_ips(clients)
    end
  rescue ApacheLogRegex::ParseError
    # nil
  end
end