IP cache

requirements.txt

@@ -39,3 +39,4 @@ plotly==4.5.0
 pdoc==0.3.2
 markdown>=3.0
 kafka-python==2.0.1
+cachetools

requirements_unit_tests.txt

@@ -34,3 +34,4 @@ isoweek==1.3.3
 pdoc==0.3.2
 spark-testing-base
 kafka-python==2.0.1
+cachetools

src/baskerville/db/models.py

@@ -115,6 +115,10 @@ class RequestSet(Base, SerializableMixin):
     process_flag = Column(Boolean, default=True)
     prediction = Column(Integer)
     attack_prediction = Column(Integer)
+    challenged = Column(Integer)
+    challenge_failed = Column(Integer)
+    challenge_passed = Column(Integer)
+    banned = Column(Integer)
     low_rate_attack = Column(Integer)
     score = Column(Float)
     features = Column(JSON)
@@ -159,6 +163,7 @@ class RequestSet(Base, SerializableMixin):
         'prediction',
         'attack_prediction',
         'low_rate_attack',
+        'challenged',
         'score',
         'label',
         'id_attribute',

src/baskerville/models/banjax_report_consumer.py

@@ -1,11 +1,15 @@
+import datetime
 import threading
 import json
 from kafka import KafkaConsumer, KafkaProducer
 import time
 import logging
 import sys
 import types
+
+from baskerville.db import set_up_db
 from baskerville.models.config import KafkaConfig
+from baskerville.models.ip_cache import IPCache
 from baskerville.util.helpers import parse_config
 import argparse
 import os
@@ -35,9 +39,12 @@ class BanjaxReportConsumer(object):
         "proxy.process.eventloop.time.max"
     ]
 
-    def __init__(self, kafka_config, logger):
-        self.config = kafka_config
+    def __init__(self, config, logger):
+        self.config = config
+        self.kafka_config = config.kafka
         self.logger = logger
+        self.ip_cache = IPCache(config, self.logger)
+        self.session, self.engine = set_up_db(config.database.__dict__)
 
         # XXX i think the metrics registry swizzling code is passing
         # an extra argument here mistakenly?.?.
@@ -49,14 +56,14 @@ def _tmp_fun(_, _2, message):
 
     def run(self):
         consumer = KafkaConsumer(
-            self.config.banjax_report_topic,
+            self.kafka_config.banjax_report_topic,
             group_id=None,
-            bootstrap_servers=self.config.bootstrap_servers,
-            security_protocol=self.config.security_protocol,
-            ssl_check_hostname=self.config.ssl_check_hostname,
-            ssl_cafile=self.config.ssl_cafile,
-            ssl_certfile=self.config.ssl_certfile,
-            ssl_keyfile=self.config.ssl_keyfile,
+            bootstrap_servers=self.kafka_config.bootstrap_servers,
+            security_protocol=self.kafka_config.security_protocol,
+            ssl_check_hostname=self.kafka_config.ssl_check_hostname,
+            ssl_cafile=self.kafka_config.ssl_cafile,
+            ssl_certfile=self.kafka_config.ssl_certfile,
+            ssl_keyfile=self.kafka_config.ssl_keyfile,
         )
 
         for message in consumer:
@@ -91,8 +98,72 @@ def consume_message(self, message):
             # 'ip_failed_challenge'-type messages are reported when a challenge is failed
             elif d.get("name") == "ip_failed_challenge":
                 self.consume_ip_failed_challenge_message(d)
+            elif d.get("name") == "ip_passed_challenge":
+                self.consume_ip_passed_challenge_message(d)
+            elif d.get("name") == "ip_banned":
+                self.consume_ip_banned_message(d)
+
+    def get_time_filter(self):
+        return (datetime.datetime.utcnow() - datetime.timedelta(
+            minutes=self.config.engine.banjax_sql_update_filter_minutes)).strftime("%Y-%m-%d %H:%M:%S %z")
 
     def consume_ip_failed_challenge_message(self, message):
+        ip = message['value_ip']
+        num_fails = self.ip_cache.ip_failed_challenge(ip)
+        if num_fails == 0:
+            return message
+
+        try:
+            if num_fails >= self.config.engine.banjax_num_fails_to_ban:
+                self.ip_cache.ip_banned(ip)
+                sql = f'update request_sets set banned = 1 where ' \
+                      f'stop > \'{self.get_time_filter()}\' and challenged = 1 and ip = \'{ip}\''
+            else:
+                sql = f'update request_sets set challenge_failed = {num_fails} where ' \
+                      f'stop > \'{self.get_time_filter()}\' and challenged = 1 and ip = \'{ip}\''
+
+            self.session.execute(sql)
+            self.session.commit()
+
+        except Exception:
+            self.session.rollback()
+            self.logger.error(Exception)
+            raise
+
+        return message
+
+    def consume_ip_passed_challenge_message(self, message):
+        ip = message['value_ip']
+        processed = self.ip_cache.ip_passed_challenge(ip)
+        if not processed:
+            return message
+        try:
+            sql = f'update request_sets set challenge_passed = 1 where ' \
+                  f'stop > \'{self.get_time_filter()}\' and challenged = 1 and ip = \'{ip}\''
+            self.session.execute(sql)
+            self.session.commit()
+
+        except Exception:
+            self.session.rollback()
+            self.logger.error(Exception)
+            raise
+
+        return message
+
+    def consume_ip_banned_message(self, message):
+        ip = message['value_ip']
+        self.logger.info(f'Banjax ip_banned {ip} ...')
+        try:
+            sql = f'update request_sets set banned = 1 where ' \
+                  f'stop > \'{self.get_time_filter()}\' and challenged = 1 and ip = \'{ip}\''
+            self.session.execute(sql)
+            self.session.commit()
+
+        except Exception:
+            self.session.rollback()
+            self.logger.error(Exception)
+            raise
+
         return message
 
 

src/baskerville/models/config.py

@@ -270,7 +270,15 @@ class EngineConfig(Config):
     sliding_window = 360
     low_rate_attack_period = [600, 3600]
     low_rate_attack_total_request = [400, 2000]
+    ip_cache_passed_challenge_ttl = 60*60*24  # 24h
+    ip_cache_passed_challenge_size = 100000
+    ip_cache_pending_ttl = 60*60*1  # 1h
+    ip_cache_pending_size = 100000
+
     white_list = None
+    banjax_sql_update_filter_minutes = 30
+    banjax_num_fails_to_ban = 9
+    register_banjax_metrics = False
 
     def __init__(self, config, parent=None):
         super(EngineConfig, self).__init__(config, parent)

src/baskerville/models/engine.py

@@ -3,9 +3,7 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
-import threading
 
-from baskerville.models.banjax_report_consumer import BanjaxReportConsumer
 from baskerville.models.base import BaskervilleBase
 from baskerville.models.config import BaskervilleConfig
 from baskerville.models.pipeline_factory import PipelineFactory
@@ -34,9 +32,7 @@ def __init__(self, run_type, conf, register_metrics=True):
         )
         self.config = BaskervilleConfig(self.config).validate()
 
-        self.register_metrics = (
-            self.config.engine.metrics and register_metrics
-        )
+        self.register_metrics = self.config.engine.metrics and register_metrics
 
         self.logger = get_logger(
             self.__class__.__name__,
@@ -216,44 +212,6 @@ def register_pipeline_metrics(self):
 
         self.logger.info('Registered metrics.')
 
-    def register_banjax_metrics(self):
-        from baskerville.util.enums import MetricClassEnum
-
-        def incr_counter_for_ip_failed_challenge(metric, self, return_value):
-            metric.labels(return_value.get('value_ip'), return_value.get('value_site')).inc()
-            return return_value
-
-        consume_ip_failed_challenge_message = metrics_registry.register_action_hook(
-            self.report_consumer.consume_ip_failed_challenge_message,
-            incr_counter_for_ip_failed_challenge,
-            metric_name='ip_failed_challenge_on_website',
-            metric_cls=MetricClassEnum.counter,
-            labelnames=['ip', 'website']
-        )
-
-        setattr(self.report_consumer, 'consume_ip_failed_challenge_message', consume_ip_failed_challenge_message)
-
-        for field_name in self.report_consumer.status_message_fields:
-            target_method = getattr(self.report_consumer, f"consume_{field_name}")
-
-            def setter_for_field(field_name_inner):
-                def label_with_id_and_set(metric, self, return_value):
-                    metric.labels(return_value.get('id')).set(return_value.get(field_name_inner))
-                    return return_value
-
-                return label_with_id_and_set
-
-            patched_method = metrics_registry.register_action_hook(
-                target_method,
-                setter_for_field(field_name),
-                metric_name=field_name.replace('.', '_'),
-                metric_cls=MetricClassEnum.gauge,
-                labelnames=['banjax_id']
-            )
-
-            setattr(self.report_consumer, f"consume_{field_name}", patched_method)
-            self.logger.info(f"Registered metric for {field_name}")
-
     def run(self) -> None:
         """
         Run steps:
@@ -266,14 +224,6 @@ def run(self) -> None:
         self.pipeline = self._set_up_pipeline()
         self.pipeline.initialize()
 
-        # self._register_metrics()
-
-        if self.register_metrics:
-            self.report_consumer = BanjaxReportConsumer(self.config.kafka, self.logger)
-            self.register_banjax_metrics()
-            self.banjax_thread = threading.Thread(target=self.report_consumer.run)
-            self.banjax_thread.start()
-
         self.pipeline.run()
 
     def finish_up(self):
@@ -286,10 +236,6 @@ def finish_up(self):
         if self.pipeline:
             self.pipeline.finish_up()
 
-        if self.banjax_thread:
-            self.banjax_thread.kill()
-            self.banjax_thread.join()
-
         self.logger.info('{} says \'Goodbye\'.'.format(
             self.__class__.__name__
         )

src/baskerville/models/ip_cache.py

@@ -0,0 +1,112 @@
+# Copyright (c) 2020, eQualit.ie inc.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import pickle
+import threading
+
+from cachetools import TTLCache
+
+from baskerville.util.helpers import get_default_ip_cache_path
+from baskerville.util.singleton_thread_safe import SingletonThreadSafe
+
+
+class IPCache(metaclass=SingletonThreadSafe):
+
+    def __init__(self, config, logger):
+        super().__init__()
+
+        self.logger = logger
+        self.lock = threading.Lock()
+
+        folder_path = get_default_ip_cache_path()
+        if not os.path.exists(folder_path):
+            os.mkdir(folder_path)
+
+        self.full_path_passed_challenge = os.path.join(folder_path, 'ip_cache_passed_challenge.bin')
+        if os.path.exists(self.full_path_passed_challenge):
+            self.logger.info(f'Loading passed challenge IP cache from file {self.full_path_passed_challenge}...')
+            with open(self.full_path_passed_challenge, 'rb') as f:
+                self.cache_passed = pickle.load(f)
+        else:
+            self.cache_passed = TTLCache(
+                maxsize=config.engine.ip_cache_passed_challenge_size,
+                ttl=config.engine.ip_cache_passed_challenge_ttl)
+            self.logger.info('A new instance of passed challege IP cache has been created')
+
+        self.full_path_pending = os.path.join(folder_path, 'ip_cache_pending.bin')
+        if os.path.exists(self.full_path_pending):
+            self.logger.info(f'Loading pending challenge IP cache from file {self.full_path_pending}...')
+            with open(self.full_path_pending, 'rb') as f:
+                self.cache_pending = pickle.load(f)
+        else:
+            self.cache_pending = TTLCache(
+                maxsize=config.engine.ip_cache_pending_size,
+                ttl=config.engine.ip_cache_pending_ttl)
+            self.logger.info('A new instance of pending IP cache has been created')
+
+    def update(self, records):
+        with self.lock:
+            self.logger.info('IP cache updating...')
+            if len(self.cache_passed) > 0.98 * self.cache_passed.maxsize:
+                self.logger.warning('IP cache passed challenge is 98% full. ')
+            if len(self.cache_pending) > 0.98 * self.cache_pending.maxsize:
+                self.logger.warning('IP cache pending challenge is 98% full. ')
+            result = []
+            for r in records:
+                if r['ip'] not in self.cache_passed and r['ip'] not in self.cache_pending:
+                    result.append(r)
+
+            for r in result:
+                self.cache_pending[r['ip']] = {
+                    'fails': 0
+                }
+
+            with open(self.full_path_pending, 'wb') as f:
+                pickle.dump(self.cache_pending, f)
+            self.logger.info(f'IP cache pending: {len(self.cache_pending)}, {len(result)} added')
+
+            return result
+
+    def ip_failed_challenge(self, ip):
+        with self.lock:
+            if ip not in self.cache_pending.keys():
+                return 0
+
+            try:
+                value = self.cache_pending[ip]
+                value['fails'] += 1
+                num_fails = value['fails']
+                self.cache_pending['ip'] = value
+                return num_fails
+
+            except KeyError as er:
+                self.logger.info(f'IP cache key error {er}')
+                pass
+
+    def ip_passed_challenge(self, ip):
+        with self.lock:
+            if ip in self.cache_passed.keys():
+                return False
+            if ip not in self.cache_pending.keys():
+                return False
+            self.cache_passed[ip] = self.cache_pending[ip]
+            del self.cache_pending[ip]
+            self.logger.info(f'IP {ip} passed challenge. Total IP in cache_passed: {len(self.cache_passed)}')
+
+            with open(self.full_path_passed_challenge, 'wb') as f:
+                pickle.dump(self.cache_passed, f)
+            self.logger.info(f'IP cache passed: {len(self.cache_passed)}, 1 added')
+        return True
+
+    def ip_banned(self, ip):
+        with self.lock:
+            try:
+                del self.cache_pending[ip]
+
+            except KeyError as er:
+                self.logger.info(f'IP cache key error {er}')
+                pass