From bcb26b38b82684ca233606052fc55004d4741971 Mon Sep 17 00:00:00 2001
From: Stan Hu <stanhu@gmail.com>
Date: Mon, 11 May 2015 15:00:54 -0700
Subject: [PATCH] Support the ability to reset Fail2Ban count and ban flag

Closes #113
---
 lib/rack/attack/cache.rb                   | 22 ++++++++++++++++++----
 lib/rack/attack/fail2ban.rb                | 15 +++++++++++----
 spec/fail2ban_spec.rb                      | 20 ++++++++++++++++++++
 spec/integration/rack_attack_cache_spec.rb | 20 ++++++++++++++++++++
 4 files changed, 69 insertions(+), 8 deletions(-)

diff --git a/lib/rack/attack/cache.rb b/lib/rack/attack/cache.rb
index c7ef6100..21f3e8f8 100644
--- a/lib/rack/attack/cache.rb
+++ b/lib/rack/attack/cache.rb
@@ -15,10 +15,7 @@ def store=(store)
       end
 
       def count(unprefixed_key, period)
-        epoch_time = Time.now.to_i
-        # Add 1 to expires_in to avoid timing error: http://git.io/i1PHXA
-        expires_in = period - (epoch_time % period) + 1
-        key = "#{prefix}:#{(epoch_time/period).to_i}:#{unprefixed_key}"
+        key, expires_in = key_and_expiry(unprefixed_key, period)
         do_count(key, expires_in)
       end
 
@@ -30,7 +27,24 @@ def write(unprefixed_key, value, expires_in)
         store.write("#{prefix}:#{unprefixed_key}", value, :expires_in => expires_in)
       end
 
+      def reset_count(unprefixed_key, period)
+        key, _ = key_and_expiry(unprefixed_key, period)
+        store.delete(key)
+      end
+
+      def delete(unprefixed_key)
+        store.delete("#{prefix}:#{unprefixed_key}")
+      end
+
       private
+
+      def key_and_expiry(unprefixed_key, period)
+        epoch_time = Time.now.to_i
+        # Add 1 to expires_in to avoid timing error: http://git.io/i1PHXA
+        expires_in = period - (epoch_time % period) + 1
+        ["#{prefix}:#{(epoch_time / period).to_i}:#{unprefixed_key}", expires_in]
+      end
+
       def do_count(key, expires_in)
         result = store.increment(key, 1, :expires_in => expires_in)
 
diff --git a/lib/rack/attack/fail2ban.rb b/lib/rack/attack/fail2ban.rb
index 443182d7..c83489a6 100644
--- a/lib/rack/attack/fail2ban.rb
+++ b/lib/rack/attack/fail2ban.rb
@@ -15,6 +15,17 @@ def filter(discriminator, options)
           end
         end
 
+        def reset(discriminator, options)
+          findtime = options[:findtime] or raise ArgumentError, "Must pass findtime option"
+          cache.reset_count("#{key_prefix}:count:#{discriminator}", findtime)
+          # Clear ban flag just in case it's there
+          cache.delete("#{key_prefix}:ban:#{discriminator}")
+        end
+
+        def banned?(discriminator)
+          cache.read("#{key_prefix}:ban:#{discriminator}") ? true : false
+        end
+
         protected
         def key_prefix
           'fail2ban'
@@ -35,10 +46,6 @@ def ban!(discriminator, bantime)
           cache.write("#{key_prefix}:ban:#{discriminator}", 1, bantime)
         end
 
-        def banned?(discriminator)
-          cache.read("#{key_prefix}:ban:#{discriminator}")
-        end
-
         def cache
           Rack::Attack.cache
         end
diff --git a/spec/fail2ban_spec.rb b/spec/fail2ban_spec.rb
index 2a4f842a..e131cf16 100644
--- a/spec/fail2ban_spec.rb
+++ b/spec/fail2ban_spec.rb
@@ -58,7 +58,27 @@
           key = "rack::attack:fail2ban:ban:1.2.3.4"
           @cache.store.read(key).must_equal 1
         end
+      end
+
+      describe 'reset after success' do
+        before do
+          get '/?test=OMGHAX', {}, 'REMOTE_ADDR' => '1.2.3.4'
+          Rack::Attack::Fail2Ban.reset('1.2.3.4', @f2b_options)
+          get '/', {}, 'REMOTE_ADDR' => '1.2.3.4'
+        end
 
+        it 'succeeds' do
+          last_response.status.must_equal 200
+        end
+
+        it 'resets fail count' do
+          key = "rack::attack:#{Time.now.to_i/@findtime}:fail2ban:count:1.2.3.4"
+          @cache.store.read(key).must_equal nil
+        end
+
+        it 'IP is not banned' do
+          Rack::Attack::Fail2Ban.banned?('1.2.3.4').must_equal false
+        end
       end
     end
   end
diff --git a/spec/integration/rack_attack_cache_spec.rb b/spec/integration/rack_attack_cache_spec.rb
index 11012b9f..c7c9df23 100644
--- a/spec/integration/rack_attack_cache_spec.rb
+++ b/spec/integration/rack_attack_cache_spec.rb
@@ -89,6 +89,26 @@ def sleep_until_expired
           @cache.read('cache-test-key').must_equal nil
         end
       end
+
+      describe "cache#delete" do
+        it "must delete the value" do
+          @cache.write("cache-test-key", "foobar", 1)
+          store.read(@key).must_equal "foobar"
+          @cache.delete('cache-test-key')
+          store.read(@key).must_be :nil?
+        end
+      end
+
+      describe "reset_count" do
+        it "must delete the value" do
+          period = 1.minute
+          unprefixed_key = 'cache-test-key'
+          @cache.count(unprefixed_key, period)
+          period_key, _ = @cache.send(:key_and_expiry, 'cache-test-key', period)
+          @cache.reset_count(unprefixed_key, period)
+          store.read(period_key).must_equal nil
+        end
+      end
     end
 
   end