Prevent out of bounds error when traversing near 0.0.0.0 (fixes #29)

fierce/fierce.py

@@ -168,12 +168,12 @@ def default_expander(ip):
 
 
 def traverse_expander(ip, n=5):
-    class_c = get_class_c_network(ip)
-
-    result = [ipaddress.IPv4Address(ip + i) for i in range(-n, n + 1)]
-    result = [i for i in result if i in class_c]
+    ip = int(ip)
+    class_c_floor = ip - (ip % 256)
 
-    return result
+    ip_min = max(ip - n, class_c_floor)
+    ip_max = min(ip + n, class_c_floor + 255)
+    return [ipaddress.IPv4Address(i) for i in range(ip_min, ip_max + 1)]
 
 
 def wide_expander(ip):

tests/test_fierce.py

@@ -27,6 +27,28 @@ def to_text(self):
         return self.response
 
 
+class TestTraverseExpander(unittest.TestCase):
+
+    def test_stay_in_class_C(self):
+        result = fierce.traverse_expander(ipaddress.IPv4Address('127.0.0.1'))
+        expected = [ipaddress.IPv4Address("127.0.0.%s" % i) for i in range(0, 7)]
+        self.assertEqual(expected, result)
+
+    # Upper and lower bound tests are to avoid reintroducing out of
+    # bounds error from IPv4Address. (Class C test won't necessarily
+    # cover this.)
+
+    def test_lower_bound(self):
+        result = fierce.traverse_expander(ipaddress.IPv4Address('0.0.0.2'), n=5)
+        expected = [ipaddress.IPv4Address(i) for i in range(0, 8)]
+        self.assertEqual(expected, result)
+
+    def test_upper_bound(self):
+        result = fierce.traverse_expander(ipaddress.IPv4Address('255.255.255.254'), n=5)
+        expected = [ipaddress.IPv4Address("255.255.255.%s" % i) for i in range(249, 256)]
+        self.assertEqual(expected, result)
+
+
 class TestFierce(unittest.TestCase):
 
     def test_concatenate_subdomains_empty(self):