diff --git a/bbot/scanner/target.py b/bbot/scanner/target.py index c321439f60..a91954557b 100644 --- a/bbot/scanner/target.py +++ b/bbot/scanner/target.py @@ -167,6 +167,21 @@ def __len__(self): def __bool__(self): return bool(len(self._rt)) or bool(self.event_seeds) + def __getstate__(self): + return { + "event_seeds": self.event_seeds, + "strict_scope": self.strict_scope, + "acl_mode": self._rt._acl_mode, + } + + def __setstate__(self, state): + self.strict_scope = state["strict_scope"] + self._rt = RadixTarget(strict_scope=state["strict_scope"], acl_mode=state["acl_mode"]) + self.event_seeds = set() + for event_seed in state["event_seeds"]: + self.event_seeds.add(event_seed) + self._add(event_seed.host, data=event_seed) + def __eq__(self, other): return self.hash == getattr(other, "hash", None) @@ -243,6 +258,10 @@ def __init__(self, *args, **kwargs): self.blacklist_regexes = set() super().__init__(*args, **kwargs) + def __setstate__(self, state): + self.blacklist_regexes = set() + super().__setstate__(state) + def get(self, host, **kwargs): """ Blacklists only accept IPs or strings. This is cleaner since we need to search for regex patterns. diff --git a/bbot/test/test_step_1/test_target.py b/bbot/test/test_step_1/test_target.py index a1344f4b30..48ab9a2d50 100644 --- a/bbot/test/test_step_1/test_target.py +++ b/bbot/test/test_step_1/test_target.py @@ -680,3 +680,33 @@ def counting_EventSeed(input): f"EventSeed was called {call_count} times for {len(targets)} targets; " f"expected {len(targets)} (seeds should reuse pre-parsed EventSeed objects)" ) + + +def test_target_pickle(): + """BBOTTarget must survive pickle round-trips (used by HTTP engine subprocess).""" + import pickle + + from bbot.scanner.target import BBOTTarget + + target = BBOTTarget( + target=["evilcorp.com", "1.2.3.0/24"], + blacklist=["bad.evilcorp.com"], + strict_scope=False, + ) + + data = pickle.dumps(target) + restored = pickle.loads(data) + + # scope checks work after unpickling + assert restored.in_target("evilcorp.com") + assert restored.in_target("www.evilcorp.com") + assert restored.in_target("1.2.3.4") + assert not restored.in_target("9.9.9.9") + + # blacklist works after unpickling + assert restored.blacklisted("bad.evilcorp.com") + assert not restored.in_scope("bad.evilcorp.com") + assert restored.in_scope("good.evilcorp.com") + + # hashes match + assert target.hash == restored.hash