bbot 2.0.1.4650rc0__tar.gz → 2.1.0.5028rc0__tar.gz

{bbot-2.0.1.4650rc0 → bbot-2.1.0.5028rc0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: bbot
-Version: 2.0.1.4650rc0
+Version: 2.1.0.5028rc0
 Summary: OSINT automation for hackers.
 Home-page: https://github.com/blacklanternsecurity/bbot
 License: GPL-3.0
@@ -14,6 +14,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: Security
 Requires-Dist: ansible (>=7.3,<9.0)
 Requires-Dist: ansible-runner (>=2.3.2,<3.0.0)
@@ -26,7 +27,7 @@ Requires-Dist: httpx (>=0.27.0,<0.28.0)
 Requires-Dist: idna (>=3.4,<4.0)
 Requires-Dist: jinja2 (>=3.1.3,<4.0.0)
 Requires-Dist: lxml (>=4.9.2,<6.0.0)
-Requires-Dist: mmh3 (>=4.1.0,<5.0.0)
+Requires-Dist: mmh3 (>=4.1,<6.0)
 Requires-Dist: omegaconf (>=2.3.0,<3.0.0)
 Requires-Dist: psutil (>=5.9.4,<6.0.0)
 Requires-Dist: pycryptodome (>=3.17,<4.0)
@@ -355,13 +356,17 @@ For more information, see [Targets](https://www.blacklanternsecurity.com/bbot/St
 
 Similar to Amass or Subfinder, BBOT supports API keys for various third-party services such as SecurityTrails, etc.
 
-The standard way to do this is to enter your API keys in **`~/.config/bbot/bbot.yml`**:
+The standard way to do this is to enter your API keys in **`~/.config/bbot/bbot.yml`**. Note that multiple API keys are allowed:
 ```yaml
 modules:
   shodan_dns:
     api_key: 4f41243847da693a4f356c0486114bc6
   c99:
-    api_key: 21a270d5f59c9b05813a72bb41707266
+    # multiple API keys
+    api_key:
+      - 21a270d5f59c9b05813a72bb41707266
+      - ea8f243d9885cf8ce9876a580224fd3c
+      - 5bc6ed268ab6488270e496d3183a1a27
   virustotal:
     api_key: dd5f0eee2e4a99b71a939bded450b246
   securitytrails:

{bbot-2.0.1.4650rc0 → bbot-2.1.0.5028rc0}/README.md

@@ -303,13 +303,17 @@ For more information, see [Targets](https://www.blacklanternsecurity.com/bbot/St
 
 Similar to Amass or Subfinder, BBOT supports API keys for various third-party services such as SecurityTrails, etc.
 
-The standard way to do this is to enter your API keys in **`~/.config/bbot/bbot.yml`**:
+The standard way to do this is to enter your API keys in **`~/.config/bbot/bbot.yml`**. Note that multiple API keys are allowed:
 ```yaml
 modules:
   shodan_dns:
     api_key: 4f41243847da693a4f356c0486114bc6
   c99:
-    api_key: 21a270d5f59c9b05813a72bb41707266
+    # multiple API keys
+    api_key:
+      - 21a270d5f59c9b05813a72bb41707266
+      - ea8f243d9885cf8ce9876a580224fd3c
+      - 5bc6ed268ab6488270e496d3183a1a27
   virustotal:
     api_key: dd5f0eee2e4a99b71a939bded450b246
   securitytrails:

{bbot-2.0.1.4650rc0 → bbot-2.1.0.5028rc0}/bbot/__init__.py

@@ -1,4 +1,4 @@
 # version placeholder (replaced by poetry-dynamic-versioning)
-__version__ = "v2.0.1.4650rc"
+__version__ = "v2.1.0.5028rc"
 
 from .scanner import Scanner, Preset

{bbot-2.0.1.4650rc0 → bbot-2.1.0.5028rc0}/bbot/core/core.py

@@ -34,8 +34,6 @@ class BBOTCore:
         self._logger = None
         self._files_config = None
 
-        self.bbot_sudo_pass = None
-
         self._config = None
         self._custom_config = None
 

{bbot-2.0.1.4650rc0 → bbot-2.1.0.5028rc0}/bbot/core/engine.py

@@ -643,6 +643,7 @@ class EngineServer(EngineBase):
                     self.log.warning(f"{self.name}: Timeout after {timeout:,} seconds in finished_tasks({tasks})")
                     for task in tasks:
                         task.cancel()
+                        self._await_cancelled_task(task)
                 else:
                     if not in_exception_chain(e, (KeyboardInterrupt, asyncio.CancelledError)):
                         self.log.error(f"{self.name}: Unhandled exception in finished_tasks({tasks}): {e}")
@@ -664,9 +665,9 @@ class EngineServer(EngineBase):
                 child_task.cancel()
 
         for task in [parent_task] + list(child_tasks):
-            await self._cancel_task(task)
+            await self._await_cancelled_task(task)
 
-    async def _cancel_task(self, task):
+    async def _await_cancelled_task(self, task):
         try:
             await asyncio.wait_for(task, timeout=10)
         except (TimeoutError, asyncio.exceptions.TimeoutError):
@@ -683,4 +684,4 @@ class EngineServer(EngineBase):
             await self.cancel_task(client_id)
         for client_id, tasks in self.child_tasks.items():
             for task in tasks:
-                await self._cancel_task(task)
+                await self._await_cancelled_task(task)

{bbot-2.0.1.4650rc0 → bbot-2.1.0.5028rc0}/bbot/core/event/base.py

@@ -1,5 +1,6 @@
 import io
 import re
+import uuid
 import json
 import base64
 import logging
@@ -24,6 +25,7 @@ from bbot.core.helpers import (
     is_domain,
     is_subdomain,
     is_ip,
+    is_ip_type,
     is_ptr,
     is_uri,
     url_depth,
@@ -58,9 +60,10 @@ class BaseEvent:
 
     Attributes:
         type (str): Specifies the type of the event, e.g., `IP_ADDRESS`, `DNS_NAME`.
-        id (str): A unique identifier for the event.
+        id (str): An identifier for the event (event type + sha1 hash of data). NOT universally unique.
+        uuid (UUID): A universally unique identifier for the event.
         data (str or dict): The main data for the event, e.g., a URL or IP address.
-        data_graph (str): Representation of `self.data` for Neo4j graph nodes.
+        data_graph (str): Representation of `self.data` for graph nodes (e.g. Neo4j).
         data_human (str): Representation of `self.data` for human output.
         data_id (str): Representation of `self.data` used to calculate the event's ID (and ultimately its hash, which is used for deduplication)
         data_json (str): Representation of `self.data` to be used in JSON serialization.
@@ -75,6 +78,7 @@ class BaseEvent:
         resolved_hosts (list of str): List of hosts to which the event data resolves, applicable for URLs and DNS names.
         parent (BaseEvent): The parent event that led to the discovery of this event.
         parent_id (str): The `id` attribute of the parent event.
+        parent_uuid (str): The `uuid` attribute of the parent event.
         tags (set of str): Descriptive tags for the event, e.g., `mx-record`, `in-scope`.
         module (BaseModule): The module that discovered the event.
         module_sequence (str): The sequence of modules that participated in the discovery.
@@ -154,7 +158,7 @@ class BaseEvent:
         Raises:
             ValidationError: If either `scan` or `parent` are not specified and `_dummy` is False.
         """
-
+        self._uuid = uuid.uuid4()
         self._id = None
         self._hash = None
         self._data = None
@@ -166,11 +170,13 @@ class BaseEvent:
         self._parent = None
         self._priority = None
         self._parent_id = None
+        self._parent_uuid = None
         self._host_original = None
         self._scope_distance = None
         self._module_priority = None
         self._resolved_hosts = set()
         self.dns_children = dict()
+        self.raw_dns_records = dict()
         self._discovery_context = ""
         self._discovery_context_regex = re.compile(r"\{(?:event|module)[^}]*\}")
         self.web_spider_distance = 0
@@ -347,6 +353,12 @@ class BaseEvent:
                 return 80
         return self._port
 
+    @property
+    def netloc(self):
+        if self.host and is_ip_type(self.host, network=False):
+            return make_netloc(self.host, self.port)
+        return None
+
     @property
     def host_stem(self):
         """
@@ -379,10 +391,20 @@ class BaseEvent:
         """
         This event's full discovery context, including those of all its parents
         """
-        parent_path = []
+        discovery_path = []
+        if self.parent is not None and self.parent is not self:
+            discovery_path = self.parent.discovery_path
+        return discovery_path + [self.discovery_context]
+
+    @property
+    def parent_chain(self):
+        """
+        This event's full discovery context, including those of all its parents
+        """
+        parent_chain = []
         if self.parent is not None and self.parent is not self:
-            parent_path = self.parent.discovery_path
-        return parent_path + [[self.id, self.discovery_context]]
+            parent_chain = self.parent.parent_chain
+        return parent_chain + [str(self.uuid)]
 
     @property
     def words(self):
@@ -425,11 +447,6 @@ class BaseEvent:
         no_host_information = not bool(self.host)
         return self._always_emit or always_emit_tags or no_host_information
 
-    @property
-    def quick_emit(self):
-        no_host_information = not bool(self.host)
-        return self._quick_emit or no_host_information
-
     @property
     def id(self):
         """
@@ -439,6 +456,13 @@ class BaseEvent:
             self._id = f"{self.type}:{self.data_hash.hex()}"
         return self._id
 
+    @property
+    def uuid(self):
+        """
+        A universally unique identifier for the event
+        """
+        return f"{self.type}:{self._uuid}"
+
     @property
     def data_hash(self):
         """
@@ -556,6 +580,13 @@ class BaseEvent:
             return parent_id
         return self._parent_id
 
+    @property
+    def parent_uuid(self):
+        parent_uuid = getattr(self.get_parent(), "uuid", None)
+        if parent_uuid is not None:
+            return parent_uuid
+        return self._parent_uuid
+
     @property
     def validators(self):
         """
@@ -722,12 +753,12 @@ class BaseEvent:
         Returns:
             dict: JSON-serializable dictionary representation of the event object.
         """
-        # type, ID, scope description
         j = dict()
-        for i in ("type", "id", "scope_description"):
+        # type, ID, scope description
+        for i in ("type", "id", "uuid", "scope_description", "netloc"):
             v = getattr(self, i, "")
             if v:
-                j.update({i: v})
+                j.update({i: str(v)})
         # event data
         data_attr = getattr(self, f"data_{mode}", None)
         if data_attr is not None:
@@ -743,6 +774,8 @@ class BaseEvent:
             j["host"] = str(self.host)
             j["resolved_hosts"] = sorted(str(h) for h in self.resolved_hosts)
             j["dns_children"] = {k: list(v) for k, v in self.dns_children.items()}
+        if isinstance(self.port, int):
+            j["port"] = self.port
         # web spider distance
         web_spider_distance = getattr(self, "web_spider_distance", None)
         if web_spider_distance is not None:
@@ -758,6 +791,9 @@ class BaseEvent:
         parent_id = self.parent_id
         if parent_id:
             j["parent"] = parent_id
+        parent_uuid = self.parent_uuid
+        if parent_uuid:
+            j["parent_uuid"] = parent_uuid
         # tags
         if self.tags:
             j.update({"tags": list(self.tags)})
@@ -770,6 +806,7 @@ class BaseEvent:
         # discovery context
         j["discovery_context"] = self.discovery_context
         j["discovery_path"] = self.discovery_path
+        j["parent_chain"] = self.parent_chain
 
         # normalize non-primitive python objects
         for k, v in list(j.items()):
@@ -909,6 +946,10 @@ class SCAN(BaseEvent):
     def discovery_path(self):
         return []
 
+    @property
+    def parent_chain(self):
+        return []
+
 
 class FINISHED(BaseEvent):
     """
@@ -1053,6 +1094,17 @@ class DnsEvent(BaseEvent):
             if parent_module_type == "DNS":
                 self.dns_resolve_distance += 1
         # self.add_tag(f"resolve-distance-{self.dns_resolve_distance}")
+        # tag subdomain / domain
+        if is_subdomain(self.host):
+            self.add_tag("subdomain")
+        elif is_domain(self.host):
+            self.add_tag("domain")
+        # tag private IP
+        try:
+            if self.host.is_private:
+                self.add_tag("private-ip")
+        except AttributeError:
+            pass
 
 
 class IP_RANGE(DnsEvent):
@@ -1069,13 +1121,6 @@ class IP_RANGE(DnsEvent):
 
 
 class DNS_NAME(DnsEvent):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        if is_subdomain(self.data):
-            self.add_tag("subdomain")
-        elif is_domain(self.data):
-            self.add_tag("domain")
-
     def sanitize_data(self, data):
         return validators.validate_host(data)
 
@@ -1498,11 +1543,18 @@ class FILESYSTEM(DictPathEvent):
     pass
 
 
-class RAW_DNS_RECORD(DictHostEvent):
+class RAW_DNS_RECORD(DictHostEvent, DnsEvent):
     # don't emit raw DNS records for affiliates
     _always_emit_tags = ["target"]
 
 
+class MOBILE_APP(DictEvent):
+    _always_emit = True
+
+    def _pretty_string(self):
+        return self.data["url"]
+
+
 def make_event(
     data,
     event_type=None,
@@ -1671,6 +1723,9 @@ def event_from_json(j, siem_friendly=False):
             data = j["data"]
         kwargs["data"] = data
         event = make_event(**kwargs)
+        event_uuid = j.get("uuid", None)
+        if event_uuid is not None:
+            event._uuid = uuid.UUID(event_uuid.split(":")[-1])
 
         resolved_hosts = j.get("resolved_hosts", [])
         event._resolved_hosts = set(resolved_hosts)
@@ -1680,6 +1735,10 @@ def event_from_json(j, siem_friendly=False):
         parent_id = j.get("parent", None)
         if parent_id is not None:
             event._parent_id = parent_id
+        parent_uuid = j.get("parent_uuid", None)
+        if parent_uuid is not None:
+            parent_type, parent_uuid = parent_uuid.split(":", 1)
+            event._parent_uuid = parent_type + ":" + str(uuid.UUID(parent_uuid))
         return event
     except KeyError as e:
         raise ValidationError(f"Event missing required field: {e}")

{bbot-2.0.1.4650rc0 → bbot-2.1.0.5028rc0}/bbot/core/helpers/command.py

@@ -295,7 +295,7 @@ def _prepare_command_kwargs(self, command, kwargs):
     if sudo and os.geteuid() != 0:
         self.depsinstaller.ensure_root()
         env["SUDO_ASKPASS"] = str((self.tools_dir / self.depsinstaller.askpass_filename).resolve())
-        env["BBOT_SUDO_PASS"] = self.depsinstaller._sudo_password
+        env["BBOT_SUDO_PASS"] = self.depsinstaller.encrypted_sudo_pw
         kwargs["env"] = env
 
         PATH = os.environ.get("PATH", "")

{bbot-2.0.1.4650rc0 → bbot-2.1.0.5028rc0}/bbot/core/helpers/depsinstaller/installer.py

@@ -10,10 +10,11 @@ from pathlib import Path
 from threading import Lock
 from itertools import chain
 from contextlib import suppress
+from secrets import token_bytes
 from ansible_runner.interface import run
 from subprocess import CalledProcessError
 
-from ..misc import can_sudo_without_password, os_platform
+from ..misc import can_sudo_without_password, os_platform, rm_at_exit
 
 log = logging.getLogger("bbot.core.helpers.depsinstaller")
 
@@ -29,14 +30,13 @@ class DepsInstaller:
         http_timeout = self.web_config.get("http_timeout", 30)
         os.environ["ANSIBLE_TIMEOUT"] = str(http_timeout)
 
+        # cache encrypted sudo pass
         self.askpass_filename = "sudo_askpass.py"
+        self._sudo_password = None
+        self._sudo_cache_setup = False
+        self._setup_sudo_cache()
         self._installed_sudo_askpass = False
-        self._sudo_password = os.environ.get("BBOT_SUDO_PASS", None)
-        if self._sudo_password is None:
-            if self.core.bbot_sudo_pass is not None:
-                self._sudo_password = self.core.bbot_sudo_pass
-            elif can_sudo_without_password():
-                self._sudo_password = ""
+
         self.data_dir = self.parent_helper.cache_dir / "depsinstaller"
         self.parent_helper.mkdir(self.data_dir)
         self.setup_status_cache = self.data_dir / "setup_status.json"
@@ -314,20 +314,27 @@ class DepsInstaller:
 
     def ensure_root(self, message=""):
         self._install_sudo_askpass()
+        # skip if we've already done this
+        if self._sudo_password is not None:
+            return
         with self.ensure_root_lock:
-            if os.geteuid() != 0 and self._sudo_password is None:
-                if message:
-                    log.warning(message)
-                while not self._sudo_password:
-                    # sleep for a split second to flush previous log messages
-                    sleep(0.1)
-                    password = getpass.getpass(prompt="[USER] Please enter sudo password: ")
-                    if self.parent_helper.verify_sudo_password(password):
-                        log.success("Authentication successful")
-                        self._sudo_password = password
-                        self.core.bbot_sudo_pass = password
-                    else:
-                        log.warning("Incorrect password")
+            # first check if the environment variable is set
+            _sudo_password = os.environ.get("BBOT_SUDO_PASS", None)
+            if _sudo_password is not None or os.geteuid() == 0 or can_sudo_without_password():
+                # if we're already root or we can sudo without a password, there's no need to prompt
+                return
+
+            if message:
+                log.warning(message)
+            while not self._sudo_password:
+                # sleep for a split second to flush previous log messages
+                sleep(0.1)
+                _sudo_password = getpass.getpass(prompt="[USER] Please enter sudo password: ")
+                if self.parent_helper.verify_sudo_password(_sudo_password):
+                    log.success("Authentication successful")
+                    self._sudo_password = _sudo_password
+                else:
+                    log.warning("Incorrect password")
 
     def install_core_deps(self):
         to_install = set()
@@ -343,6 +350,38 @@ class DepsInstaller:
             self.ensure_root()
             self.apt_install(list(to_install))
 
+    def _setup_sudo_cache(self):
+        if not self._sudo_cache_setup:
+            self._sudo_cache_setup = True
+            # write temporary encryption key, to be deleted upon scan completion
+            self._sudo_temp_keyfile = self.parent_helper.temp_filename()
+            # remove it at exit
+            rm_at_exit(self._sudo_temp_keyfile)
+            # generate random 32-byte key
+            random_key = token_bytes(32)
+            # write key to file and set secure permissions
+            self._sudo_temp_keyfile.write_bytes(random_key)
+            self._sudo_temp_keyfile.chmod(0o600)
+            # export path to environment variable, for use in askpass script
+            os.environ["BBOT_SUDO_KEYFILE"] = str(self._sudo_temp_keyfile.resolve())
+
+    @property
+    def encrypted_sudo_pw(self):
+        if self._sudo_password is None:
+            return ""
+        return self._encrypt_sudo_pw(self._sudo_password)
+
+    def _encrypt_sudo_pw(self, pw):
+        from Crypto.Cipher import AES
+        from Crypto.Util.Padding import pad
+
+        key = self._sudo_temp_keyfile.read_bytes()
+        cipher = AES.new(key, AES.MODE_CBC)
+        ct_bytes = cipher.encrypt(pad(pw.encode(), AES.block_size))
+        iv = cipher.iv.hex()
+        ct = ct_bytes.hex()
+        return f"{iv}:{ct}"
+
     def _install_sudo_askpass(self):
         if not self._installed_sudo_askpass:
             self._installed_sudo_askpass = True

bbot-2.1.0.5028rc0/bbot/core/helpers/depsinstaller/sudo_askpass.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python3
+import os
+import sys
+from pathlib import Path
+from Crypto.Cipher import AES
+from Crypto.Util.Padding import unpad
+
+ENV_VAR_NAME = "BBOT_SUDO_PASS"
+KEY_ENV_VAR_PATH = "BBOT_SUDO_KEYFILE"
+
+
+def decrypt_password(encrypted_data, key):
+    iv, ciphertext = encrypted_data.split(":")
+    iv = bytes.fromhex(iv)
+    ct = bytes.fromhex(ciphertext)
+    cipher = AES.new(key, AES.MODE_CBC, iv)
+    pt = unpad(cipher.decrypt(ct), AES.block_size)
+    return pt.decode("utf-8")
+
+
+def main():
+    encrypted_password = os.environ.get(ENV_VAR_NAME, "")
+    # remove variable from environment once we've got it
+    os.environ.pop(ENV_VAR_NAME, None)
+    encryption_keypath = Path(os.environ.get(KEY_ENV_VAR_PATH, ""))
+
+    if not encrypted_password or not encryption_keypath.is_file():
+        print("Error: Encrypted password or encryption key not found in environment variables.", file=sys.stderr)
+        sys.exit(1)
+
+    try:
+        key = encryption_keypath.read_bytes()
+        decrypted_password = decrypt_password(encrypted_password, key)
+        print(decrypted_password, end="")
+    except Exception as e:
+        print(f'Error decrypting password "{encrypted_password}": {str(e)}', file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

{bbot-2.0.1.4650rc0 → bbot-2.1.0.5028rc0}/bbot/core/helpers/dns/brute.py

@@ -15,15 +15,17 @@ class DNSBrute:
     >>> results = await self.helpers.dns.brute(self, domain, subdomains)
     """
 
-    nameservers_url = (
+    _nameservers_url = (
         "https://raw.githubusercontent.com/blacklanternsecurity/public-dns-servers/master/nameservers.txt"
     )
 
     def __init__(self, parent_helper):
         self.parent_helper = parent_helper
         self.log = logging.getLogger("bbot.helper.dns.brute")
+        self.dns_config = self.parent_helper.config.get("dns", {})
         self.num_canaries = 100
-        self.max_resolvers = self.parent_helper.config.get("dns", {}).get("brute_threads", 1000)
+        self.max_resolvers = self.dns_config.get("brute_threads", 1000)
+        self.nameservers_url = self.dns_config.get("brute_nameservers", self._nameservers_url)
         self.devops_mutations = list(self.parent_helper.word_cloud.devops_mutations)
         self.digit_regex = self.parent_helper.re.compile(r"\d+")
         self._resolver_file = None
@@ -39,18 +41,12 @@ class DNSBrute:
             type = "A"
         type = str(type).strip().upper()
 
-        domain_wildcard_rdtypes = set()
-        for _domain, rdtypes in (await self.parent_helper.dns.is_wildcard_domain(domain)).items():
-            for rdtype, results in rdtypes.items():
-                if results:
-                    domain_wildcard_rdtypes.add(rdtype)
-        if any([r in domain_wildcard_rdtypes for r in (type, "CNAME")]):
-            self.log.info(
-                f"Aborting massdns on {domain} because it's a wildcard domain ({','.join(domain_wildcard_rdtypes)})"
+        wildcard_rdtypes = await self.parent_helper.dns.is_wildcard_domain(domain, (type, "CNAME"))
+        if wildcard_rdtypes:
+            self.log.hugewarning(
+                f"Aborting massdns on {domain} because it's a wildcard domain ({','.join(wildcard_rdtypes)})"
             )
             return []
-        else:
-            self.log.debug(f"{domain}: A is not in domain_wildcard_rdtypes:{domain_wildcard_rdtypes}")
 
         canaries = self.gen_random_subdomains(self.num_canaries)
         canaries_list = list(canaries)
@@ -148,10 +144,15 @@ class DNSBrute:
 
     async def resolver_file(self):
         if self._resolver_file is None:
-            self._resolver_file = await self.parent_helper.wordlist(
+            self._resolver_file_original = await self.parent_helper.wordlist(
                 self.nameservers_url,
                 cache_hrs=24 * 7,
             )
+            nameservers = set(self.parent_helper.read_file(self._resolver_file_original))
+            nameservers.difference_update(self.parent_helper.dns.system_resolvers)
+            # exclude system nameservers from brute-force
+            # this helps prevent rate-limiting which might cause BBOT's main dns queries to fail
+            self._resolver_file = self.parent_helper.tempfile(nameservers, pipe=False)
         return self._resolver_file
 
     def gen_random_subdomains(self, n=50):

{bbot-2.0.1.4650rc0 → bbot-2.1.0.5028rc0}/bbot/core/helpers/dns/dns.py

@@ -66,7 +66,7 @@ class DNSHelper(EngineClient):
         self.resolver.timeout = self.timeout
         self.resolver.lifetime = self.timeout
 
-        self.runaway_limit = self.config.get("runaway_limit", 5)
+        self.runaway_limit = self.dns_config.get("runaway_limit", 5)
 
         # wildcard handling
         self.wildcard_disable = self.dns_config.get("wildcard_disable", False)
@@ -117,8 +117,11 @@ class DNSHelper(EngineClient):
             self._brute = DNSBrute(self.parent_helper)
         return self._brute
 
-    @async_cachedmethod(lambda self: self._is_wildcard_cache)
-    async def is_wildcard(self, query, ips=None, rdtype=None):
+    @async_cachedmethod(
+        lambda self: self._is_wildcard_cache,
+        key=lambda query, rdtypes, raw_dns_records: (query, tuple(sorted(rdtypes)), bool(raw_dns_records)),
+    )
+    async def is_wildcard(self, query, rdtypes, raw_dns_records=None):
         """
         Use this method to check whether a *host* is a wildcard entry
 
@@ -150,9 +153,6 @@ class DNSHelper(EngineClient):
         Note:
             `is_wildcard` can be True, False, or None (indicating that wildcard detection was inconclusive)
         """
-        if [ips, rdtype].count(None) == 1:
-            raise ValueError("Both ips and rdtype must be specified")
-
         query = self._wildcard_prevalidation(query)
         if not query:
             return {}
@@ -161,15 +161,17 @@ class DNSHelper(EngineClient):
         if is_domain(query):
             return {}
 
-        return await self.run_and_return("is_wildcard", query=query, ips=ips, rdtype=rdtype)
+        return await self.run_and_return("is_wildcard", query=query, rdtypes=rdtypes, raw_dns_records=raw_dns_records)
 
-    @async_cachedmethod(lambda self: self._is_wildcard_domain_cache)
-    async def is_wildcard_domain(self, domain, log_info=False):
+    @async_cachedmethod(
+        lambda self: self._is_wildcard_domain_cache, key=lambda domain, rdtypes: (domain, tuple(sorted(rdtypes)))
+    )
+    async def is_wildcard_domain(self, domain, rdtypes):
         domain = self._wildcard_prevalidation(domain)
         if not domain:
             return {}
 
-        return await self.run_and_return("is_wildcard_domain", domain=domain, log_info=False)
+        return await self.run_and_return("is_wildcard_domain", domain=domain, rdtypes=rdtypes)
 
     def _wildcard_prevalidation(self, host):
         if self.wildcard_disable:
@@ -192,8 +194,8 @@ class DNSHelper(EngineClient):
 
         return host
 
-    async def _mock_dns(self, mock_data):
+    async def _mock_dns(self, mock_data, custom_lookup_fn=None):
         from .mock import MockResolver
 
-        self.resolver = MockResolver(mock_data)
-        await self.run_and_return("_mock_dns", mock_data=mock_data)
+        self.resolver = MockResolver(mock_data, custom_lookup_fn=custom_lookup_fn)
+        await self.run_and_return("_mock_dns", mock_data=mock_data, custom_lookup_fn=custom_lookup_fn)