Pull filter and expander functionality out of subdomain enumeration loop

fierce/fierce.py

@@ -158,6 +158,10 @@ def get_class_c_network(ip):
     return class_c
 
 
+def default_expander(ip):
+    return [ip]
+
+
 def traverse_expander(ip, n=5):
     class_c = get_class_c_network(ip)
 
@@ -175,11 +179,26 @@ def wide_expander(ip):
     return result
 
 
+def range_expander(ip):
+    network = ipaddress.IPv4Network(ip)
+
+    result = list(network)
+
+    return result
+
+
+def default_filter(address):
+    return True
+
+
 def search_filter(domains, address):
     return any(domain in address for domain in domains)
 
 
 def find_nearby(resolver, ips, filter_func=None):
+    if filter_func is None:
+        filter_func = default_filter
+
     str_ips = [str(ip) for ip in ips]
 
     # https://docs.python.org/3.5/library/concurrent.futures.html#concurrent.futures.ThreadPoolExecutor
@@ -198,10 +217,11 @@ def find_nearby(resolver, ips, filter_func=None):
             )
         }
 
-    reversed_ips = {k: v for k, v in reversed_ips.items() if v is not None}
-
-    if filter_func:
-        reversed_ips = {k: v for k, v in reversed_ips.items() if v and filter_func(v[0].to_text())}
+    reversed_ips = {
+        k: v
+        for k, v in reversed_ips.items()
+        if v is not None and filter_func(v[0].to_text())
+    }
 
     return reversed_ips
 
@@ -256,8 +276,11 @@ def fierce(**kwargs):
     )
 
     if kwargs.get("range"):
-        internal_range = ipaddress.IPv4Network(kwargs.get("range"))
-        nearby_ips = find_nearby(resolver, list(internal_range))
+        range_ips = range_expander(kwargs.get("range"))
+        nearby_ips = find_nearby(
+            resolver,
+            range_ips,
+        )
         if nearby_ips:
             print("Nearby:")
             pprint.pprint({k: v[0].to_text() for k, v in nearby_ips.items() if v})
@@ -305,7 +328,17 @@ def fierce(**kwargs):
         kwargs["subdomain_file"]
     )
 
-    get_unvisited = unvisited_closure()
+    filter_func = None
+    if kwargs.get("search"):
+        filter_func = functools.partial(search_filter, kwargs["search"])
+
+    expander_func = default_expander
+    if kwargs.get("wide"):
+        expander_func = wide_expander
+    elif kwargs.get("traverse"):
+        expander_func = functools.partial(traverse_expander, n=kwargs["traverse"])
+
+    unvisited = unvisited_closure()
 
     for subdomain in subdomains:
         url = concatenate_subdomains(domain, [subdomain])
@@ -320,20 +353,14 @@ def fierce(**kwargs):
         if kwargs.get('connect') and not ip.is_private:
             http_connection_headers = head_request(str(ip))
 
-        filter_func = None
-        if kwargs.get("search"):
-            filter_func = functools.partial(search_filter, kwargs["search"])
-
-        if kwargs.get("wide"):
-            ips = wide_expander(ip)
-        elif kwargs.get("traverse"):
-            ips = traverse_expander(ip, kwargs["traverse"])
-        else:
-            ips = []
+        ips = expander_func(ip)
+        unvisited_ips = unvisited(ips)
 
-        ips = get_unvisited(ips)
-
-        nearby_ips = find_nearby(resolver, ips, filter_func=filter_func)
+        nearby_ips = find_nearby(
+            resolver,
+            unvisited_ips,
+            filter_func=filter_func
+        )
 
         nearby = None
         if nearby_ips:

tests/test_fierce.py

@@ -13,6 +13,17 @@
 from fierce import fierce
 
 
+# Simply getting a dns.resolver.Answer with a specific result was
+# more difficult than I'd like, let's just go with this less than
+# ideal approach for now
+class MockAnswer(object):
+    def __init__(self, response):
+        self.response = response
+
+    def to_text(self):
+        return self.response
+
+
 class TestFierce(unittest.TestCase):
 
     def test_concatenate_subdomains_empty(self):
@@ -79,6 +90,16 @@ def test_concatenate_subdomains_fqdn_subdomain(self):
 
         self.assertEqual(expected, result)
 
+    def test_default_expander(self):
+        ip = ipaddress.IPv4Address('192.168.1.1')
+
+        result = fierce.default_expander(ip)
+        expected = [
+            ipaddress.IPv4Address('192.168.1.1'),
+        ]
+
+        self.assertEqual(expected, result)
+
     def test_traverse_expander_basic(self):
         ip = ipaddress.IPv4Address('192.168.1.1')
         expand = 1
@@ -156,6 +177,18 @@ def test_wide_expander_upper_boundary(self):
 
         self.assertEqual(expected, result)
 
+    def test_range_expander(self):
+        ip = '192.168.1.0/31'
+
+        result = fierce.range_expander(ip)
+
+        expected = [
+            ipaddress.IPv4Address('192.168.1.0'),
+            ipaddress.IPv4Address('192.168.1.1'),
+        ]
+
+        self.assertEqual(expected, result)
+
     def test_recursive_query_basic_failure(self):
         resolver = dns.resolver.Resolver()
         domain = dns.name.from_text('example.com.')
@@ -292,17 +325,19 @@ def test_find_nearby_basic(self):
             ipaddress.IPv4Address('192.168.1.0'),
             ipaddress.IPv4Address('192.168.1.1'),
         ]
+        returned_answer1 = [MockAnswer('sd1.example.com.')]
+        returned_answer2 = [MockAnswer('sd2.example.com.')]
         side_effect = [
-            'sd1.example.com.',
-            'sd2.example.com.',
+            returned_answer1,
+            returned_answer2,
         ]
 
         with unittest.mock.patch.object(fierce, 'reverse_query', side_effect=side_effect):
             result = fierce.find_nearby(resolver, ips)
 
         expected = {
-            '192.168.1.0': 'sd1.example.com.',
-            '192.168.1.1': 'sd2.example.com.',
+            '192.168.1.0': returned_answer1,
+            '192.168.1.1': returned_answer2,
         }
 
         self.assertEqual(expected, result)
@@ -313,19 +348,7 @@ def test_find_nearby_filter_func(self):
             ipaddress.IPv4Address('192.168.1.0'),
             ipaddress.IPv4Address('192.168.1.1'),
         ]
-
-        # Simply getting a dns.resolver.Answer with a specific result was
-        # more difficult than I'd like, let's just go with this less than
-        # ideal approach for now
-        class MockAnswer(object):
-            def __init__(self, response):
-                self.response = response
-
-            def to_text(self):
-                return self.response
-
         returned_answer = [MockAnswer('sd1.example.com.')]
-
         side_effect = [
             returned_answer,
             [MockAnswer('sd2.example.com.')],
@@ -447,6 +470,30 @@ def test_unvisited_closure_overlapping_intersection(self):
 
         self.assertEqual(expected, result)
 
+    def test_search_filter_empty(self):
+        domains = []
+        address = 'test.example.com'
+
+        result = fierce.search_filter(domains, address)
+
+        self.assertFalse(result)
+
+    def test_search_filter_true(self):
+        domains = ['example.com']
+        address = 'test.example.com'
+
+        result = fierce.search_filter(domains, address)
+
+        self.assertTrue(result)
+
+    def test_search_filter_false(self):
+        domains = ['not.com']
+        address = 'test.example.com'
+
+        result = fierce.search_filter(domains, address)
+
+        self.assertFalse(result)
+
 
 if __name__ == "__main__":
     unittest.main()