Fix/limits

csirtg_smrt/rule.py

@@ -34,6 +34,7 @@ def __init__(self, path=None, rule=None, **kwargs):
                 self.filters = d.get('filters')
                 self.delim_pattern = d.get('delim_pattern')
                 self.line_filter = d.get('line_filter')
+                self.limit = d.get('limit')
 
             else:
                 raise RuleUnsupported('unsupported file type: {}'.format(path))
@@ -55,6 +56,7 @@ def __init__(self, path=None, rule=None, **kwargs):
             self.filters = rule.get('filters')
             self.delim_pattern = rule.get('delim_pattern')
             self.line_filter = rule.get('line_filter')
+            self.limit = rule.get('limit')
 
     def __repr__(self):
         return json.dumps({
@@ -70,5 +72,6 @@ def __repr__(self):
             'itype': self.itype,
             'filters': self.filters,
             'delim_pattern': self.delim_pattern,
-            'line_filter': self.line_filter
+            'line_filter': self.line_filter,
+            'limit': self.limit
         }, sort_keys=True, indent=4, separators=(',', ': '))

csirtg_smrt/smrt.py

@@ -216,6 +216,10 @@ def process(self, rule, feed, limit=None, data=None, filters=None):
         parser = self.load_parser(rule, feed, limit=limit, data=data, filters=filters)
 
         feed_indicators = parser.process()
+
+        if not limit:
+            limit = rule.feeds[feed].get('limit')
+
         if limit:
             feed_indicators = itertools.islice(feed_indicators, int(limit))
 

test/alexa/alexa.yml

@@ -23,6 +23,7 @@ defaults:
     95)))
 
 feeds:
-  top1m:
+  topN:
     remote: http://s3.amazonaws.com/alexa-static/top-1m.csv.zip
+    limit: 10100
 

test/alexa/test_alexa.py

@@ -13,8 +13,7 @@
 
 
 def test_alexa_top1m():
-    print(rule)
-    r, feed = next(s.load_feeds(rule, feed='top1m'))
+    r, feed = next(s.load_feeds(rule, feed='topN'))
 
     r.defaults['remote'] = 'test/alexa/alexa_top-1m.csv.zip'
     x = s.process(r, feed)
@@ -33,6 +32,8 @@ def test_alexa_top1m():
     assert int(x[1000].confidence) == 25
     assert int(x[9999].confidence) == 25
     assert int(x[10000].confidence) == 0
+
+    assert len(x) == 10100
 
     tags = set()
     for xx in x:

test/csirtg/test_csirtg.py

@@ -84,3 +84,12 @@ def test_csirtg_skips_first():
     x = s.process(rule, feed="port-scanners")
     x = list(x)
     assert len(x) == 3
+
+
+def test_csirtg_limits():
+    rule.feeds['port-scanners']['remote'] = 'test/csirtg/feed2_csv.txt'
+    rule.feeds['port-scanners']['limit'] = 1
+
+    x = s.process(rule, feed="port-scanners")
+    x = list(x)
+    assert len(x) == 1