# BruteForceKilla
#
#   A Rack middleware to limit requests by ip address, coded for fun as my first
#   middleware, thanks http://coderack.org for giving me a reason :)
#
#   For production use, one would want to make a memcache or redis tracker.
#
#   options:
#
#     :tracker      => Class name of the tracker to use (default Memory (all there is for now!))
#     :sample_count => # of requests from ip before calcing the duration (default 10)
#     :duration     => Max allowed duration in seconds (default 1 second)
 
module Rack
  module BruteForceKillaTracker
    # Tracks ip addresses in memory, local to the process.  Not recommeded for most
    # apps, provided as a contract example, specifically trackers need to:
    #
    # 1) initialize the :sample_count, :duration options (other opts pass-thru from middleware)
    # 2) respond_to kill?(env) -- you can grab the ip address from env['REMOTE_ADDR']
    class Memory
      
      def initialize(options={})
        @sample_count = options[:sample_count] 
        @duration = options[:duration]
      end
      
      def ip_addresses
        @ip_addresses ||= {}
      end
      
      def kill?(env)
        ip_address = ['REMOTE_ADDR']
        hits = (ip_addresses[ip_address] ||= [])
        hits << Time.now.to_i
        if hits.size >= @sample_count
          sum = 0
          for x in 0..(@sample_count-2)
            sum += (hits[x+1] - hits[x])
          end
          ip_addresses.delete(ip_address)
          (sum/@sample_count) < @duration ? true : false
        else
          ip_addresses.clear if ip_addresses.size > 2000
          false
        end
      end
    end
  end
end
 
module Rack
  class BruteForceKilla
 
    def initialize(app,options={})
      @tracker = (options.delete(:tracker) || BruteForceKillaTracker::Memory).new({:sample_count=>10,:duration=>1}.merge!(options))
      @app = app
    end
    
    def call(env)
      if @tracker.kill?(env)
        msg = "Request rate violation"
        [403,{"Content-Length"=>msg.length.to_s},msg]
      else
        @app.call(env)
      end
    end
 
  end
end