From 6ee2bc21a9e2e65d19f9437dbdf1bcd8c41c8e36 Mon Sep 17 00:00:00 2001
From: Wes Young <wes@barely3am.com>
Date: Wed, 15 Feb 2017 11:23:24 -0500
Subject: [PATCH 1/2] fixes #141

---
 csirtg_smrt/rule.py | 5 ++++-
 csirtg_smrt/smrt.py | 4 ++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/csirtg_smrt/rule.py b/csirtg_smrt/rule.py
index f8c3a09..5d58a28 100644
--- a/csirtg_smrt/rule.py
+++ b/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=(',', ': '))
diff --git a/csirtg_smrt/smrt.py b/csirtg_smrt/smrt.py
index 0e5c5ce..e326e4f 100644
--- a/csirtg_smrt/smrt.py
+++ b/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))
 

From ed20306a5697530e9b17ee491693b411993d9609 Mon Sep 17 00:00:00 2001
From: Wes Young <wes@barely3am.com>
Date: Wed, 15 Feb 2017 11:32:03 -0500
Subject: [PATCH 2/2] fixing tests

---
 test/alexa/alexa.yml       | 3 ++-
 test/alexa/test_alexa.py   | 5 +++--
 test/csirtg/test_csirtg.py | 9 +++++++++
 3 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/test/alexa/alexa.yml b/test/alexa/alexa.yml
index b016b24..16d45c4 100644
--- a/test/alexa/alexa.yml
+++ b/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
     
\ No newline at end of file
diff --git a/test/alexa/test_alexa.py b/test/alexa/test_alexa.py
index a596acf..93347a9 100644
--- a/test/alexa/test_alexa.py
+++ b/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:
diff --git a/test/csirtg/test_csirtg.py b/test/csirtg/test_csirtg.py
index beb9882..ac89ef2 100644
--- a/test/csirtg/test_csirtg.py
+++ b/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