Better cloud tagging

bbot/modules/internal/cloudcheck.py

@@ -1,3 +1,5 @@
+from contextlib import suppress
+
 from bbot.modules.base import BaseInterceptModule
 
 
@@ -28,15 +30,28 @@ async def handle_event(self, event, **kwargs):
         if self.dummy_modules is None:
             self.make_dummy_modules()
         # cloud tagging by hosts
-        hosts_to_check = set(str(s) for s in event.resolved_hosts)
-        # we use the original host, since storage buckets hostnames might be collapsed to _wildcard
-        hosts_to_check.add(str(event.host_original))
-        for host in hosts_to_check:
+        hosts_to_check = set(event.resolved_hosts)
+        with suppress(KeyError):
+            hosts_to_check.remove(event.host_original)
+        hosts_to_check = [event.host_original] + list(hosts_to_check)
+
+        for i, host in enumerate(hosts_to_check):
+            host_is_ip = self.helpers.is_ip(host)
             for provider, provider_type, subnet in self.helpers.cloudcheck(host):
                 if provider:
                     event.add_tag(f"{provider_type}-{provider}")
+                    if host_is_ip:
+                        event.add_tag(f"{provider_type}-ip")
+                    else:
+                        # if the original hostname is a cloud domain, tag it as such
+                        if i == 0:
+                            event.add_tag(f"{provider_type}-domain")
+                        # any children are tagged as CNAMEs
+                        else:
+                            event.add_tag(f"{provider_type}-cname")
 
         found = set()
+        str_hosts_to_check = [str(host) for host in hosts_to_check]
         # look for cloud assets in hosts, http responses
         # loop through each provider
         for provider in self.helpers.cloud.providers.values():
@@ -54,7 +69,7 @@ async def handle_event(self, event, **kwargs):
                     if event.type == "HTTP_RESPONSE":
                         matches = await self.helpers.re.findall(sig, event.data.get("body", ""))
                     elif event.type.startswith("DNS_NAME"):
-                        for host in hosts_to_check:
+                        for host in str_hosts_to_check:
                             match = sig.match(host)
                             if match:
                                 matches.append(match.groups())

bbot/test/test_step_2/module_tests/test_module_cloudcheck.py

@@ -51,6 +51,10 @@ async def setup_after_prep(self, module_test):
             await module.handle_event(event)
             assert "cloud-amazon" in event.tags, f"{event} was not properly cloud-tagged"
 
+        assert "cloud-domain" in aws_event1.tags
+        assert "cloud-ip" in other_event2.tags
+        assert "cloud-cname" in other_event3.tags
+
         for event in (aws_event3, other_event1):
             await module.handle_event(event)
             assert "cloud-amazon" not in event.tags, f"{event} was improperly cloud-tagged"

bbot/test/test_step_2/module_tests/test_module_portscan.py

@@ -109,10 +109,12 @@ def check(self, module_test, events):
                 if e.type == "DNS_NAME" and e.data == "dummy.asdf.evilcorp.net" and str(e.module) == "dummy_module"
             ]
         )
-        assert 2 <= len([e for e in events if e.type == "IP_ADDRESS" and e.data == "8.8.8.8"]) <= 3
-        assert 2 <= len([e for e in events if e.type == "IP_ADDRESS" and e.data == "8.8.4.4"]) <= 3
-        assert 2 <= len([e for e in events if e.type == "IP_ADDRESS" and e.data == "8.8.4.5"]) <= 3
-        assert 2 <= len([e for e in events if e.type == "IP_ADDRESS" and e.data == "8.8.4.6"]) <= 3
+        # the reason these numbers aren't exactly predictable is because we can't predict which one arrives first
+        # to the portscan module. Sometimes, one that would normally be deduped is force-emitted because it led to a new open port.
+        assert 2 <= len([e for e in events if e.type == "IP_ADDRESS" and e.data == "8.8.8.8"]) <= 4
+        assert 2 <= len([e for e in events if e.type == "IP_ADDRESS" and e.data == "8.8.4.4"]) <= 4
+        assert 2 <= len([e for e in events if e.type == "IP_ADDRESS" and e.data == "8.8.4.5"]) <= 4
+        assert 2 <= len([e for e in events if e.type == "IP_ADDRESS" and e.data == "8.8.4.6"]) <= 4
         assert 1 == len([e for e in events if e.type == "OPEN_TCP_PORT" and e.data == "8.8.8.8:443"])
         assert 1 == len([e for e in events if e.type == "OPEN_TCP_PORT" and e.data == "8.8.4.5:80"])
         assert 1 == len([e for e in events if e.type == "OPEN_TCP_PORT" and e.data == "8.8.4.6:631"])