Set a limit on how many alerts can be merged (airbnb#727)

* Re-add valid_ip helper

* Limit the max number of alerts which can be merged

helpers/base.py

@@ -23,7 +23,8 @@
     # Import some utility functions which are useful for rules as well
     get_first_key,
     get_keys,
-    in_network
+    in_network,
+    valid_ip
 )
 
 logging.basicConfig()

stream_alert/alert_merger/main.py

@@ -28,6 +28,9 @@
 
 class AlertMergeGroup(object):
     """A list of alerts within a single merge window which match on their merge keys."""
+    # In order to limit the size of a merged alert, cap the maximum number that can be merged.
+    MAX_ALERTS_PER_GROUP = 50
+
     def __init__(self, alert):
         """Initialize the group with the oldest alert remaining."""
         self.alerts = [alert]
@@ -38,6 +41,9 @@ def add(self, alert):
         Returns:
             True if the alert matches this group and was added, False otherwise.
         """
+        if len(self.alerts) >= self.MAX_ALERTS_PER_GROUP:
+            return False
+
         if alert.can_merge(self.alerts[0]):
             self.alerts.append(alert)
             return True

tests/unit/stream_alert_alert_merger/test_main.py

@@ -163,6 +163,22 @@ def test_merge_groups_complex(self):
         assert_equal(alerts[5:7], groups[2].alerts)
         assert_equal([alerts[7]], groups[3].alerts)
 
+    @patch.object(main.AlertMergeGroup, 'MAX_ALERTS_PER_GROUP', 2)
+    def test_merge_groups_limit_reached(self):
+        """Alert Merger - Alert Collection - Max Alerts Per Group"""
+        alerts = [
+            Alert('same_rule_name', {'key': 'A'}, set(),
+                  created=datetime(year=2000, month=1, day=1),
+                  merge_by_keys=['key'], merge_window=timedelta(minutes=5)),
+        ] * 5
+
+        # Since max alerts per group is 2, it should create 3 merged groups.
+        groups = main.AlertMerger._merge_groups(alerts)
+        assert_equal(3, len(groups))
+        assert_equal(alerts[0:2], groups[0].alerts)
+        assert_equal(alerts[2:4], groups[1].alerts)
+        assert_equal(alerts[4:], groups[2].alerts)
+
     @patch.object(main, 'LOGGER')
     @patch.object(main.AlertMerger, 'MAX_LAMBDA_PAYLOAD_SIZE', 600)
     def test_dispatch(self, mock_logger):