wkimeria-rack-attack 4.1.2

data/lib/rack/attack/allow2ban.rb

@@ -0,0 +1,23 @@
+module Rack
+  class 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

data/lib/rack/attack/blacklist.rb

@@ -0,0 +1,12 @@
+module Rack
+  class Attack
+    class Blacklist < Check
+      def initialize(name, block)
+        super
+        @type = :blacklist
+      end
+
+    end
+  end
+end
+

data/lib/rack/attack/cache.rb

@@ -0,0 +1,57 @@
+module Rack
+  class Attack
+    class Cache
+
+      attr_accessor :prefix
+
+      def initialize
+        self.store = ::Rails.cache if defined?(::Rails.cache)
+        @prefix = 'rack::attack'
+      end
+
+      attr_reader :store
+      def store=(store)
+        @store = StoreProxy.build(store)
+      end
+
+      def get_count(unprefixed_key, period)
+        if store.class == Rack::Attack::StoreProxy::RedisStoreProxy
+          store.raw_read(build_key(unprefixed_key, period))
+        else
+          store.read(build_key(unprefixed_key, period))
+        end
+      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
+        do_count(build_key(unprefixed_key, period), expires_in)
+      end
+
+      def read(unprefixed_key)
+        store.read("#{prefix}:#{unprefixed_key}")
+      end
+
+      def write(unprefixed_key, value, expires_in)
+        store.write("#{prefix}:#{unprefixed_key}", value, :expires_in => expires_in)
+      end
+
+      private
+      def do_count(key, expires_in)
+        result = store.increment(key, 1, :expires_in => expires_in)
+
+        # NB: Some stores return nil when incrementing uninitialized values
+        if result.nil?
+          store.write(key, 1, :expires_in => expires_in)
+        end
+        result || 1
+      end
+
+      def build_key(unprefixed_key, period)
+        epoch_time = Time.now.to_i
+        "#{prefix}:#{(epoch_time/period).to_i}:#{unprefixed_key}"
+      end
+    end
+  end
+end

data/lib/rack/attack/check.rb

@@ -0,0 +1,23 @@
+module Rack
+  class Attack
+    class Check
+      attr_reader :name, :block, :type
+      def initialize(name, options = {}, block)
+        @name, @block = name, block
+        @type = options.fetch(:type, nil)
+      end
+
+      def [](req)
+        block[req].tap {|match|
+          if match
+            req.env["rack.attack.matched"] = name
+            req.env["rack.attack.match_type"] = type
+            Rack::Attack.instrument(req)
+          end
+        }
+      end
+
+    end
+  end
+end
+

data/lib/rack/attack/conditional_throttle.rb

@@ -0,0 +1,17 @@
+module Rack
+  class Attack
+    class ConditionalThrottle < ::Rack::Attack::Throttle
+
+      def increment_counter(discriminator)
+        key = "#{name}:#{discriminator}"
+        cache.count(key, period)
+      end
+
+      def get_count(discriminator)
+        key = "#{name}:#{discriminator}"
+        count = cache.get_count(key, period)
+        count ? count.to_i : 0
+      end
+    end
+  end
+end

data/lib/rack/attack/fail2ban.rb

@@ -0,0 +1,48 @@
+module Rack
+  class Attack
+    class Fail2Ban
+      class << self
+        def filter(discriminator, options)
+          bantime   = options[:bantime]   or raise ArgumentError, "Must pass bantime option"
+          findtime  = options[:findtime]  or raise ArgumentError, "Must pass findtime option"
+          maxretry  = options[:maxretry]  or raise ArgumentError, "Must pass maxretry option"
+
+          if banned?(discriminator)
+            # Return true for blacklist
+            true
+          elsif yield
+            fail!(discriminator, bantime, findtime, maxretry)
+          end
+        end
+
+        protected
+        def key_prefix
+          'fail2ban'
+        end
+
+        def fail!(discriminator, bantime, findtime, maxretry)
+          count = cache.count("#{key_prefix}:count:#{discriminator}", findtime)
+          if count >= maxretry
+            ban!(discriminator, bantime)
+          end
+
+          true
+        end
+
+
+        private
+        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
+      end
+    end
+  end
+end

data/lib/rack/attack/request.rb

@@ -0,0 +1,19 @@
+# Rack::Attack::Request is the same as ::Rack::Request by default.
+#
+# This is a safe place to add custom helper methods to the request object
+# through monkey patching:
+#
+#   class Rack::Attack::Request < ::Rack::Request
+#     def localhost?
+#       ip == "127.0.0.1"
+#     end
+#   end
+#
+#   Rack::Attack.whitelist("localhost") {|req| req.localhost? }
+#
+module Rack
+  class Attack
+    class Request < ::Rack::Request
+    end
+  end
+end

data/lib/rack/attack/store_proxy.rb

@@ -0,0 +1,22 @@
+module Rack
+  class Attack
+    module StoreProxy
+      PROXIES = [DalliProxy, RedisStoreProxy]
+
+      def self.build(store)
+        # RedisStore#increment needs different behavior, so detect that
+        # (method has an arity of 2; must call #expire separately
+        if defined?(::ActiveSupport::Cache::RedisStore) && store.is_a?(::ActiveSupport::Cache::RedisStore)
+          # ActiveSupport::Cache::RedisStore doesn't expose any way to set an expiry,
+          # so use the raw Redis::Store instead
+          store = store.instance_variable_get(:@data)
+        end
+
+        klass = PROXIES.find { |proxy| proxy.handle?(store) }
+
+        klass ? klass.new(store) : store
+      end
+
+    end
+  end
+end

data/lib/rack/attack/store_proxy/dalli_proxy.rb

@@ -0,0 +1,65 @@
+require 'delegate'
+
+module Rack
+  class Attack
+    module StoreProxy
+      class DalliProxy < SimpleDelegator
+        def self.handle?(store)
+          return false unless defined?(::Dalli)
+
+          # Consider extracting to a separate Connection Pool proxy to reduce
+          # code here and handle clients other than Dalli.
+          if defined?(::ConnectionPool) && store.is_a?(::ConnectionPool)
+            store.with { |conn| conn.is_a?(::Dalli::Client) }
+          else
+            store.is_a?(::Dalli::Client)
+          end
+        end
+
+        def initialize(client)
+          super(client)
+          stub_with_if_missing
+        end
+
+        def read(key)
+          with do |client|
+            client.get(key)
+          end
+        rescue Dalli::DalliError
+        end
+
+        def write(key, value, options={})
+          with do |client|
+            client.set(key, value, options.fetch(:expires_in, 0), raw: true)
+          end
+        rescue Dalli::DalliError
+        end
+
+        def increment(key, amount, options={})
+          with do |client|
+            client.incr(key, amount, options.fetch(:expires_in, 0), amount)
+          end
+        rescue Dalli::DalliError
+        end
+
+        def delete(key)
+          with do |client|
+            client.delete(key)
+          end
+        rescue Dalli::DalliError
+        end
+
+        private
+
+        def stub_with_if_missing
+          unless __getobj__.respond_to?(:with)
+            class << self
+              def with; yield __getobj__; end
+            end
+          end
+        end
+
+      end
+    end
+  end
+end

data/lib/rack/attack/store_proxy/redis_store_proxy.rb

@@ -0,0 +1,49 @@
+require 'delegate'
+
+module Rack
+  class Attack
+    module StoreProxy
+      class RedisStoreProxy < SimpleDelegator
+        def self.handle?(store)
+          defined?(::Redis::Store) && store.is_a?(::Redis::Store)
+        end
+
+        def initialize(store)
+          super(store)
+        end
+
+        def raw_read(key)
+          #https://github.com/redis-store/redis-store/issues/96
+          # if raw not specified results in 'marshal data too short' error
+          self.get(key, raw: true)
+        rescue Redis::BaseError
+        end
+
+        def read(key)
+          self.get(key)
+        rescue Redis::BaseError
+        end
+
+        def write(key, value, options={})
+          if (expires_in = options[:expires_in])
+            self.setex(key, expires_in, value)
+          else
+            self.set(key, value)
+          end
+        rescue Redis::BaseError
+        end
+
+        def increment(key, amount, options={})
+          count = nil
+          self.pipelined do
+            count = self.incrby(key, amount)
+            self.expire(key, options[:expires_in]) if options[:expires_in]
+          end
+          count.value if count
+        rescue Redis::BaseError
+        end
+
+      end
+    end
+  end
+end

data/lib/rack/attack/throttle.rb

@@ -0,0 +1,50 @@
+module Rack
+  class Attack
+    class Throttle
+      MANDATORY_OPTIONS = [:limit, :period]
+      attr_reader :name, :limit, :period, :block, :type
+      def initialize(name, options, block)
+        @name, @block = name, block
+        MANDATORY_OPTIONS.each do |opt|
+          raise ArgumentError.new("Must pass #{opt.inspect} option") unless options[opt]
+        end
+        @limit  = options[:limit]
+        @period = options[:period].to_i
+        @type   = options.fetch(:type, :throttle)
+      end
+
+      def cache
+        Rack::Attack.cache
+      end
+
+      def get_count(discriminator)
+        key = "#{name}:#{discriminator}"
+        cache.count(key, period)
+      end
+
+      def [](req)
+        discriminator = block[req]
+        return false unless discriminator
+
+        count = get_count(discriminator)
+        current_limit = limit.respond_to?(:call) ? limit.call(req) : limit
+        data = {
+          :count => count,
+          :period => period,
+          :limit => current_limit
+        }
+        (req.env['rack.attack.throttle_data'] ||= {})[name] = data
+
+        (count > current_limit).tap do |throttled|
+          if throttled
+            req.env['rack.attack.matched']             = name
+            req.env['rack.attack.match_discriminator'] = discriminator
+            req.env['rack.attack.match_type']          = type
+            req.env['rack.attack.match_data']          = data
+            Rack::Attack.instrument(req)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/rack/attack/track.rb

@@ -0,0 +1,21 @@
+module Rack
+  class Attack
+    class Track
+      extend Forwardable
+
+      attr_reader :filter
+
+      def initialize(name, options = {}, block)
+        options[:type] = :track
+
+        if options[:limit] && options[:period]
+          @filter = Throttle.new(name, options, block)
+        else
+          @filter = Check.new(name, options, block)
+        end
+      end
+
+      def_delegator :@filter, :[]
+    end
+  end
+end

data/lib/rack/attack/version.rb

@@ -0,0 +1,5 @@
+module Rack
+  class Attack
+    VERSION = '4.1.2'
+  end
+end

data/lib/rack/attack/whitelist.rb

@@ -0,0 +1,11 @@
+module Rack
+  class Attack
+    class Whitelist < Check
+      def initialize(name, block)
+        super
+        @type = :whitelist
+      end
+
+    end
+  end
+end

data/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 403
+      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 403
+      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