Merge pull request #33 from teslamotors/allow2ban

Allow2Ban

README.md

@@ -128,6 +128,24 @@ how the parameters work.
     end
 ```
 
+#### Allow2Ban
+`Allow2Ban.filter` works the same way as the `Fail2Ban.filter` except that it *allows* requests from misbehaving
+clients until such time as they reach maxretry at which they are cut off as per normal.
+```ruby
+    # Lockout IP addresses that are hammering your login page.
+    # After 20 requests in 1 minute, block all requests from that IP for 1 hour.
+    Rack::Attack.blacklist('allow2ban login scrapers') do |req|
+      # `filter` returns false value if request is to your login page (but still
+      # increments the count) so request below the limit are not blocked until
+      # they hit the limit.  At that point, filter will return true and block.
+      Rack::Attack::Fail2Ban.filter(req.ip, :maxretry => 20, :findtime => 1.minute, :bantime => 1.hour) do
+        # The count for the IP is incremented if the return value is truthy.
+        req.path = '/login' and req.method == 'post'
+      end
+    end
+```
+
+
 ### Throttles
 
 ```ruby

lib/rack/attack.rb

@@ -8,6 +8,7 @@ module Rack::Attack
   autoload :Track,     'rack/attack/track'
   autoload :StoreProxy,'rack/attack/store_proxy'
   autoload :Fail2Ban,  'rack/attack/fail2ban'
+  autoload :Allow2Ban,  'rack/attack/allow2ban'
 
   class << self
 

lib/rack/attack/allow2ban.rb

@@ -0,0 +1,23 @@
+module Rack
+  module Attack
+    class Allow2Ban < Fail2Ban
+      class << self
+        protected
+        def key_prefix
+          'allow2ban'
+        end
+
+        # everything the same here except we return only return true
+        # (blocking the request) if they have tripped the limit.
+        def fail!(discriminator, bantime, findtime, maxretry)
+          count = cache.count("#{key_prefix}:count:#{discriminator}", findtime)
+          if count >= maxretry
+            ban!(discriminator, bantime)
+          end
+          # we may not block them this time, but they're banned for next time
+          false
+        end
+      end
+    end
+  end
+end

lib/rack/attack/fail2ban.rb

@@ -15,23 +15,28 @@ module Rack
           end
         end
 
-        private
+        protected
+        def key_prefix
+          'fail2ban'
+        end
+
         def fail!(discriminator, bantime, findtime, maxretry)
-          count = cache.count("fail2ban:count:#{discriminator}", findtime)
+          count = cache.count("#{key_prefix}:count:#{discriminator}", findtime)
           if count >= maxretry
             ban!(discriminator, bantime)
           end
 
-          # Return true for blacklist
           true
         end
 
+
+        private
         def ban!(discriminator, bantime)
-          cache.write("fail2ban:ban:#{discriminator}", 1, bantime)
+          cache.write("#{key_prefix}:ban:#{discriminator}", 1, bantime)
         end
 
         def banned?(discriminator)
-          cache.read("fail2ban:ban:#{discriminator}")
+          cache.read("#{key_prefix}:ban:#{discriminator}")
         end
 
         def cache

spec/allow2ban_spec.rb

@@ -0,0 +1,121 @@
+require_relative 'spec_helper'
+describe 'Rack::Attack.Allow2Ban' do
+  before do
+    # Use a long findtime; failures due to cache key rotation less likely
+    @cache = Rack::Attack.cache
+    @findtime = 60
+    @bantime  = 60
+    Rack::Attack.cache.store = ActiveSupport::Cache::MemoryStore.new
+    @f2b_options = {:bantime => @bantime, :findtime => @findtime, :maxretry => 2}
+    Rack::Attack.blacklist('pentest') do |req|
+      Rack::Attack::Allow2Ban.filter(req.ip, @f2b_options){req.query_string =~ /OMGHAX/}
+    end
+  end
+
+  describe 'discriminator has not been banned' do
+    describe 'making ok request' do
+      it 'succeeds' do
+        get '/', {}, 'REMOTE_ADDR' => '1.2.3.4'
+        last_response.status.must_equal 200
+      end
+    end
+
+    describe 'making qualifying request' do
+      describe 'when not at maxretry' do
+        before { get '/?foo=OMGHAX', {}, 'REMOTE_ADDR' => '1.2.3.4' }
+        it 'succeeds' do
+          last_response.status.must_equal 200
+        end
+
+        it 'increases fail count' do
+          key = "rack::attack:#{Time.now.to_i/@findtime}:allow2ban:count:1.2.3.4"
+          @cache.store.read(key).must_equal 1
+        end
+
+        it 'is not banned' do
+          key = "rack::attack:allow2ban:1.2.3.4"
+          @cache.store.read(key).must_be_nil
+        end
+      end
+
+      describe 'when at maxretry' do
+        before do
+          # maxretry is 2 - so hit with an extra failed request first
+          get '/?test=OMGHAX', {}, 'REMOTE_ADDR' => '1.2.3.4'
+          get '/?foo=OMGHAX', {}, 'REMOTE_ADDR' => '1.2.3.4'
+        end
+
+        it 'succeeds' do
+          last_response.status.must_equal 200
+        end
+
+        it 'increases fail count' do
+          key = "rack::attack:#{Time.now.to_i/@findtime}:allow2ban:count:1.2.3.4"
+          @cache.store.read(key).must_equal 2
+        end
+
+        it 'is banned' do
+          key = "rack::attack:allow2ban:ban:1.2.3.4"
+          @cache.store.read(key).must_equal 1
+        end
+
+      end
+    end
+  end
+
+  describe 'discriminator has been banned' do
+    before do
+      # maxretry is 2 - so hit enough times to get banned
+      get '/?test=OMGHAX', {}, 'REMOTE_ADDR' => '1.2.3.4'
+      get '/?foo=OMGHAX', {}, 'REMOTE_ADDR' => '1.2.3.4'
+    end
+
+    describe 'making request for other discriminator' do
+      it 'succeeds' do
+        get '/', {}, 'REMOTE_ADDR' => '2.2.3.4'
+        last_response.status.must_equal 200
+      end
+    end
+
+    describe 'making ok request' do
+      before do
+        get '/', {}, 'REMOTE_ADDR' => '1.2.3.4'
+      end
+
+      it 'fails' do
+        last_response.status.must_equal 401
+      end
+
+      it 'does not increase fail count' do
+        key = "rack::attack:#{Time.now.to_i/@findtime}:allow2ban:count:1.2.3.4"
+        @cache.store.read(key).must_equal 2
+      end
+
+      it 'is still banned' do
+        key = "rack::attack:allow2ban:ban:1.2.3.4"
+        @cache.store.read(key).must_equal 1
+      end
+    end
+
+    describe 'making failing request' do
+      before do
+        get '/?foo=OMGHAX', {}, 'REMOTE_ADDR' => '1.2.3.4'
+      end
+
+      it 'fails' do
+        last_response.status.must_equal 401
+      end
+
+      it 'does not increase fail count' do
+        key = "rack::attack:#{Time.now.to_i/@findtime}:allow2ban:count:1.2.3.4"
+        @cache.store.read(key).must_equal 2
+      end
+
+      it 'is still banned' do
+        key = "rack::attack:allow2ban:ban:1.2.3.4"
+        @cache.store.read(key).must_equal 1
+      end
+    end
+
+  end
+end