From 9aba2eb3a563a39a7356ff0739e79d02a9842857 Mon Sep 17 00:00:00 2001
From: William Ballenthin <william.ballenthin@fireeye.com>
Date: Tue, 30 Jun 2020 00:44:22 -0600
Subject: [PATCH] rules: range: correct handling of range with min==0

closes #57
---
 capa/engine.py       |   8 +++-
 tests/test_engine.py | 100 ++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 105 insertions(+), 3 deletions(-)

diff --git a/capa/engine.py b/capa/engine.py
index 4be1e32c..0eb164f1 100644
--- a/capa/engine.py
+++ b/capa/engine.py
@@ -154,10 +154,14 @@ class Range(Statement):
         self.max = max if max is not None else (1 << 64 - 1)
 
     def evaluate(self, ctx):
-        if self.child not in ctx:
+        count = len(ctx.get(self.child, []))
+        if self.min == 0:
+            if count == 0:
+                return Result(True, self, [])
+        elif self.child not in ctx:
+            # self.min > 0 so there needs to be more than zero matches
             return Result(False, self, [])
 
-        count = len(ctx[self.child])
         return Result(self.min <= count <= self.max, self, [], locations=ctx[self.child])
 
     def __str__(self):
diff --git a/tests/test_engine.py b/tests/test_engine.py
index 5c7c9a3c..5681ba64 100644
--- a/tests/test_engine.py
+++ b/tests/test_engine.py
@@ -65,7 +65,8 @@ def test_complex():
 
 def test_range():
     # unbounded range, but no matching feature
-    assert Range(Number(1)).evaluate({Number(2): {}}) == False
+    # since the lower bound is zero, and there are zero matches, ok
+    assert Range(Number(1)).evaluate({Number(2): {}}) == True
 
     # unbounded range with matching feature should always match
     assert Range(Number(1)).evaluate({Number(1): {}}) == True
@@ -96,6 +97,103 @@ def test_range():
     assert Range(Number(1), min=1, max=3).evaluate({Number(1): {1, 2, 3, 4}}) == False
 
 
+def test_range_exact():
+    rule = textwrap.dedent('''
+        rule:
+            meta:
+                name: test rule
+            features:
+                - count(number(100)): 2
+    ''')
+    r = capa.rules.Rule.from_yaml(rule)
+
+    # just enough matches
+    features, matches = capa.engine.match([r], {capa.features.insn.Number(100): {1, 2}}, 0x0)
+    assert 'test rule' in matches
+
+    # not enough matches
+    features, matches = capa.engine.match([r], {capa.features.insn.Number(100): {1}}, 0x0)
+    assert 'test rule' not in matches
+
+    # too many matches
+    features, matches = capa.engine.match([r], {capa.features.insn.Number(100): {1, 2, 3}}, 0x0)
+    assert 'test rule' not in matches
+
+
+def test_range_range():
+    rule = textwrap.dedent('''
+         rule:
+             meta:
+                 name: test rule
+             features:
+                 - count(number(100)): (2, 3)
+     ''')
+    r = capa.rules.Rule.from_yaml(rule)
+
+    # just enough matches
+    features, matches = capa.engine.match([r], {capa.features.insn.Number(100): {1, 2}}, 0x0)
+    assert 'test rule' in matches
+
+    # enough matches
+    features, matches = capa.engine.match([r], {capa.features.insn.Number(100): {1, 2, 3}}, 0x0)
+    assert 'test rule' in matches
+
+    # not enough matches
+    features, matches = capa.engine.match([r], {capa.features.insn.Number(100): {1}}, 0x0)
+    assert 'test rule' not in matches
+
+    # too many matches
+    features, matches = capa.engine.match([r], {capa.features.insn.Number(100): {1, 2, 3, 4}}, 0x0)
+    assert 'test rule' not in matches
+
+
+def test_range_exact_zero():
+    rule = textwrap.dedent('''
+        rule:
+            meta:
+                name: test rule
+            features:
+                - count(number(100)): 0
+    ''')
+    r = capa.rules.Rule.from_yaml(rule)
+
+    # feature isn't indexed - good.
+    features, matches = capa.engine.match([r], {}, 0x0)
+    assert 'test rule' in matches
+
+    # feature is indexed, but no matches.
+    # i don't think we should ever really have this case, but good to check anyways.
+    features, matches = capa.engine.match([r], {capa.features.insn.Number(100): {}}, 0x0)
+    assert 'test rule' in matches
+
+    # too many matches
+    features, matches = capa.engine.match([r], {capa.features.insn.Number(100): {1}}, 0x0)
+    assert 'test rule' not in matches
+
+
+def test_range_with_zero():
+    rule = textwrap.dedent('''
+         rule:
+             meta:
+                 name: test rule
+             features:
+                 - count(number(100)): (0, 1)
+     ''')
+    r = capa.rules.Rule.from_yaml(rule)
+
+    # ok
+    features, matches = capa.engine.match([r], {}, 0x0)
+    assert 'test rule' in matches
+    features, matches = capa.engine.match([r], {capa.features.insn.Number(100): {}}, 0x0)
+    assert 'test rule' in matches
+    features, matches = capa.engine.match([r], {capa.features.insn.Number(100): {1}}, 0x0)
+    assert 'test rule' in matches
+
+    # too many matches
+    features, matches = capa.engine.match([r], {capa.features.insn.Number(100): {1, 2}}, 0x0)
+    assert 'test rule' not in matches
+
+
 def test_match_adds_matched_rule_feature():
     '''show that using `match` adds a feature for matched rules.'''
     rule = textwrap.dedent('''