From 9d60cb7da90ca8a1c9936b532b0bb3d66e653bdf Mon Sep 17 00:00:00 2001
From: TheTechromancer <thetechromancer@protonmail.com>
Date: Tue, 17 Oct 2023 16:17:27 -0400
Subject: [PATCH] better scope tagging for events

---
 bbot/core/event/base.py                       | 14 ++++++-----
 .../test_manager_scope_accuracy.py            | 24 +++++++++++++++++++
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/bbot/core/event/base.py b/bbot/core/event/base.py
index e2643a026..54aa83897 100644
--- a/bbot/core/event/base.py
+++ b/bbot/core/event/base.py
@@ -349,10 +349,6 @@ def scope_distance(self, scope_distance):
             The method will automatically update the relevant 'distance-' tags associated with the event.
         """
         if scope_distance >= 0:
-            if scope_distance == 0:
-                self.add_tag("in-scope")
-            else:
-                self.remove_tag("in-scope")
             new_scope_distance = None
             # ensure scope distance does not increase (only allow setting to smaller values)
             if self.scope_distance == -1:
@@ -360,11 +356,17 @@ def scope_distance(self, scope_distance):
             else:
                 new_scope_distance = min(self.scope_distance, scope_distance)
             if self._scope_distance != new_scope_distance:
-                self._scope_distance = new_scope_distance
+                # remove old scope distance tags
                 for t in list(self.tags):
                     if t.startswith("distance-"):
                         self.remove_tag(t)
-                self.add_tag(f"distance-{new_scope_distance}")
+                if scope_distance == 0:
+                    self.add_tag("in-scope")
+                    self.remove_tag("affiliate")
+                else:
+                    self.remove_tag("in-scope")
+                    self.add_tag(f"distance-{new_scope_distance}")
+                self._scope_distance = new_scope_distance
             # apply recursively to parent events
             source_scope_distance = getattr(self.source, "scope_distance", -1)
             if source_scope_distance >= 0 and self != self.source:
diff --git a/bbot/test/test_step_1/test_manager_scope_accuracy.py b/bbot/test/test_step_1/test_manager_scope_accuracy.py
index 9b9dd0579..7022de8a3 100644
--- a/bbot/test/test_step_1/test_manager_scope_accuracy.py
+++ b/bbot/test/test_step_1/test_manager_scope_accuracy.py
@@ -790,3 +790,27 @@ async def test_manager_blacklist(bbot_config, bbot_scanner, bbot_httpserver, cap
     assert not any([e for e in events if e.type == "URL_UNVERIFIED" and e.data == "http://www-prod.test.notreal:8888/"])
 
     assert 'Omitting due to blacklisted DNS associations: URL_UNVERIFIED("http://www-prod.test.notreal:8888/"' in caplog.text
+
+
+@pytest.mark.asyncio
+async def test_manager_scope_tagging(bbot_config, bbot_scanner):
+    scan = bbot_scanner("test.notreal", config=bbot_config)
+    e1 = scan.make_event("www.test.notreal", source=scan.root_event, tags=["affiliate"])
+    log.critical(e1)
+    assert e1.scope_distance == 1
+    assert "distance-1" in e1.tags
+    assert "affiliate" in e1.tags
+
+    e2 = scan.make_event("dev.test.notreal", source=e1, tags=["affiliate"])
+    assert e2.scope_distance == 2
+    assert "in-scope" not in e2.tags
+    distance_tags = [t for t in e2.tags if t.startswith("distance-")]
+    assert len(distance_tags) == 1
+    assert distance_tags[0] == "distance-2"
+
+    e2.scope_distance = 0
+    log.critical(e2)
+    assert e2.scope_distance == 0
+    assert "in-scope" in e2.tags
+    distance_tags = [t for t in e2.tags if t.startswith("distance-")]
+    assert not distance_tags