fix: detect wildcard DNS up front and filter false positives

README.md

@@ -33,12 +33,13 @@ Seed subwiz with these subdomains:
 usage: cli.py [-h] -i INPUT_FILE [-o OUTPUT_FILE] [-n NUM_PREDICTIONS] [--no-resolve]
               [--force-download] [--max-recursion MAX_RECURSION] [-t TEMPERATURE]
               [-d {auto,cpu,cuda,mps}] [-m MAX_NEW_TOKENS]
-              [--resolution-concurrency RESOLUTION_CONCURRENCY] [--multi-apex] [-q] [-s]
+              [--resolution-concurrency RESOLUTION_CONCURRENCY] [--multi-apex]
+              [--wildcard {filter,skip}] [-q] [-s]
 
 options:
   -h, --help            show this help message and exit
   -i, --input-file INPUT_FILE
-                        file containing new-line-separated subdomains. (default: None)
+                        file containing new-line-separated subdomains.
   -o, --output-file OUTPUT_FILE
                         output file to write new-line separated subdomains to. (default: None)
   -n, --num-predictions NUM_PREDICTIONS
@@ -58,6 +59,10 @@ options:
                         number of concurrent resolutions. (default: 128)
   --multi-apex          allow multiple apex domains in the input file. runs inference for each
                         apex separately. (default: False)
+  --wildcard {filter,skip}
+                        how to handle apexes with a wildcard DNS record: 'filter' keeps only
+                        subdomains resolving outside the wildcard IP set; 'skip' drops the apex.
+                        (default: filter)
   -q, --quiet           useful for piping into another tool. (default: False)
   -s, --silent          do not print any output. requires --output-file. (default: False)
 ```

subwiz/cli.py

@@ -99,6 +99,14 @@
     dest="multi_apex",
     action="store_true",
 )
+parser.add_argument(
+    "--wildcard",
+    help="how to handle apexes with a wildcard DNS record: 'filter' keeps only "
+    "subdomains resolving outside the wildcard IP set; 'skip' drops the apex.",
+    dest="wildcard",
+    default="filter",
+    choices=["filter", "skip"],
+)
 parser.add_argument(
     "-q",
     "--quiet",

subwiz/main.py

@@ -5,6 +5,8 @@
 tokenization, inference execution, and result processing.
 """
 
+from __future__ import annotations
+
 import argparse
 import asyncio
 import os
@@ -19,7 +21,7 @@
 
 from subwiz.cli_printer import print_hello, print_log, print_progress_dot
 from subwiz.model import GPT
-from subwiz.resolve import get_registered_domains
+from subwiz.resolve import detect_wildcard, get_registered_domains
 from subwiz.type import (
     Domain,
     input_domains_type,
@@ -30,7 +32,6 @@
     concurrency_type,
 )
 
-
 MODEL_REPO = "HadrianSecurity/subwiz"
 MODEL_FILE = "model_v2.pt"
 TOKENIZER_FILE = "tokenizer_v2.json"
@@ -170,6 +171,7 @@ def _get_domains_for_group(
     no_resolve: bool,
     resolution_concurrency: int,
     quiet: bool,
+    wildcard_ips: set[str] | None = None,
 ) -> set[str]:
     """For a group of subdomains that share an apex: run inference and check if they resolve, recursively.
 
@@ -186,6 +188,8 @@ def _get_domains_for_group(
         no_resolve: Whether to skip DNS resolution
         resolution_concurrency: Number of concurrent DNS resolutions
         print_cli_progress: Whether to print progress information
+        wildcard_ips: Catch-all IPs of a detected wildcard record for this apex,
+            used to filter out false positives. None if no wildcard is present.
 
     Returns:
         Set of discovered subdomain strings
@@ -233,7 +237,7 @@ def _counting_progress_dot():
             return {str(dom) for dom in predictions}
 
         predictions_that_resolve = asyncio.run(
-            get_registered_domains(predictions, resolution_concurrency)
+            get_registered_domains(predictions, resolution_concurrency, wildcard_ips)
         )
 
         if not quiet:
@@ -272,6 +276,7 @@ def run(
     multi_apex: bool = False,
     max_recursion: int = 5,
     quiet: bool = True,
+    wildcard: str = "filter",
 ) -> list[str]:
     """Check types, download model, get new subdomains for each apex.
 
@@ -287,6 +292,9 @@ def run(
         multi_apex: Whether to allow multiple apex domains
         max_recursion: Maximum recursion depth for discovery
         print_cli_progress: Whether to print progress information
+        wildcard: How to handle apexes with a wildcard DNS record. "filter"
+            (default) keeps only subdomains resolving outside the wildcard IP
+            set; "skip" drops the apex entirely.
 
     Returns:
         List of discovered subdomain strings
@@ -305,6 +313,11 @@ def run(
     temperature = temperature_type(temperature)
     resolution_concurrency = concurrency_type(resolution_concurrency)
 
+    if wildcard not in ("filter", "skip"):
+        raise argparse.ArgumentTypeError(
+            f'wildcard should be "filter" or "skip": {wildcard}'
+        )
+
     domain_groups = defaultdict(set)
     for dom in domain_objects:
         domain_groups[dom.apex_domain].add(dom)
@@ -321,6 +334,24 @@ def run(
     found_domains = set()
 
     for apex in sorted(domain_groups):
+        wildcard_ips = None
+        if not no_resolve:
+            wildcard_ips = asyncio.run(detect_wildcard(apex, resolution_concurrency))
+            if wildcard_ips:
+                if wildcard == "skip":
+                    if not quiet:
+                        print_log(
+                            f"wildcard DNS detected on {apex}, skipping "
+                            f"(resolves to {', '.join(sorted(wildcard_ips))})"
+                        )
+                    continue
+                if not quiet:
+                    print_log(
+                        f"wildcard DNS detected on {apex}, filtering to "
+                        f"subdomains resolving outside "
+                        f"{', '.join(sorted(wildcard_ips))}"
+                    )
+
         found_domains |= _get_domains_for_group(
             domains_in_group=domain_groups[apex],
             all_apexes=set(domain_groups.keys()),
@@ -334,6 +365,7 @@ def run(
             no_resolve=no_resolve,
             resolution_concurrency=resolution_concurrency,
             quiet=quiet,
+            wildcard_ips=wildcard_ips,
         )
 
     return sorted(found_domains)

subwiz/model.py

@@ -23,7 +23,6 @@
 from transformers import PreTrainedTokenizerFast
 from typing import Callable, Optional
 
-
 VALID_SUBDOMAIN_RE = re.compile(
     r"^(?!-)([a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?)(?:\.([a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?))*(?<!\.)$",
     re.IGNORECASE,

subwiz/resolve.py

@@ -3,41 +3,131 @@
 This module provides asynchronous DNS resolution functionality to check whether
 domains are registered and resolve to IP addresses. It uses multiple nameservers
 for reliability and implements concurrency control for efficient batch processing.
+
+It also provides wildcard-DNS detection. Domains with a wildcard record
+(``*.example.com``) resolve *every* possible subdomain to the same catch-all
+IP set, so plain "does it resolve?" checks produce huge numbers of false
+positives. ``detect_wildcard`` finds the catch-all IP set up front so that
+resolution can filter to only the subdomains that resolve *outside* it.
 """
 
+from __future__ import annotations
+
 import asyncio
+import random
+import string
 
 import aiodns
 import idna.core
 
 from subwiz.type import Domain
 
-
 NAME_SERVERS = ["1.1.1.1", "1.0.0.1", "8.8.8.8"]
 TIMEOUT = 3
 TRIES = 1
 
+WILDCARD_PROBE_COUNT = 3
+WILDCARD_LABEL_LENGTH = 20
+
+
+def _new_resolver() -> aiodns.DNSResolver:
+    """Create a DNS resolver configured with the module nameservers."""
+    return aiodns.DNSResolver(nameservers=NAME_SERVERS, timeout=TIMEOUT, tries=TRIES)
+
+
+async def _resolve_ips(
+    hostname: str, resolver: aiodns.DNSResolver, semaphore: asyncio.Semaphore
+) -> set[str]:
+    """Resolve a hostname to its set of IPv4 addresses.
+
+    Args:
+        hostname: Hostname to resolve
+        resolver: DNS resolver instance to use for queries
+        semaphore: Semaphore for controlling concurrency
+
+    Returns:
+        Set of IP address strings the hostname resolves to, or an empty set if
+        it does not resolve.
+    """
+    async with semaphore:
+        try:
+            results = await resolver.query(hostname, "A")
+            return {record.host for record in results}
+        except idna.IDNAError:
+            return set()
+        except aiodns.error.DNSError:
+            return set()
+
+
+def _random_label() -> str:
+    """Generate a random subdomain label unlikely to exist as a real record."""
+    alphabet = string.ascii_lowercase + string.digits
+    return "".join(random.choices(alphabet, k=WILDCARD_LABEL_LENGTH))
+
+
+async def detect_wildcard(
+    apex_domain: str, resolution_concurrency: int = WILDCARD_PROBE_COUNT
+) -> set[str]:
+    """Detect a wildcard DNS record on an apex domain.
+
+    Probes several random subdomains that should not exist. If they all
+    resolve, the apex has a wildcard (catch-all) record and the union of the
+    IPs they resolve to is returned as the wildcard IP set.
+
+    Args:
+        apex_domain: Apex domain to probe (e.g. ``example.com``)
+        resolution_concurrency: Maximum number of concurrent DNS resolutions
+
+    Returns:
+        The set of catch-all IPs if a wildcard is present, otherwise an empty
+        set.
+    """
+    semaphore = asyncio.Semaphore(resolution_concurrency)
+    resolver = _new_resolver()
+
+    probes = [f"{_random_label()}.{apex_domain}" for _ in range(WILDCARD_PROBE_COUNT)]
+    probe_ips = await asyncio.gather(
+        *[_resolve_ips(probe, resolver, semaphore) for probe in probes]
+    )
+
+    if not all(probe_ips):
+        return set()
+
+    return set().union(*probe_ips)
+
 
 async def get_registered_domains(
-    domains_to_check: set[Domain], resolution_concurrency: int
+    domains_to_check: set[Domain],
+    resolution_concurrency: int,
+    wildcard_ips: set[str] | None = None,
 ) -> set[Domain]:
     """Check which domains from a set are registered and resolve to IP addresses.
 
+    When ``wildcard_ips`` is provided (the apex has a wildcard record), a
+    domain counts as registered only if it resolves to at least one IP
+    *outside* the wildcard set. Domains resolving solely to the catch-all IPs
+    are treated as false positives and dropped.
+
     Args:
         domains_to_check: Set of Domain objects to check for registration
         resolution_concurrency: Maximum number of concurrent DNS resolutions
+        wildcard_ips: Catch-all IPs of a detected wildcard, or None/empty for
+            the standard "resolves = registered" check.
 
     Returns:
         Set of Domain objects that are registered and resolve successfully
     """
 
     semaphore = asyncio.Semaphore(resolution_concurrency)
-    resolver = aiodns.DNSResolver(
-        nameservers=NAME_SERVERS, timeout=TIMEOUT, tries=TRIES
-    )
+    resolver = _new_resolver()
 
     domains_list = list(domains_to_check)
-    tasks = [dom.is_registered(resolver, semaphore) for dom in domains_to_check]
-    results = await asyncio.gather(*tasks)
 
-    return {dom for dom, is_reg in zip(domains_list, results) if is_reg}
+    if not wildcard_ips:
+        tasks = [dom.is_registered(resolver, semaphore) for dom in domains_list]
+        results = await asyncio.gather(*tasks)
+        return {dom for dom, is_reg in zip(domains_list, results) if is_reg}
+
+    tasks = [_resolve_ips(str(dom), resolver, semaphore) for dom in domains_list]
+    results = await asyncio.gather(*tasks)
+    return {dom for dom, ips in zip(domains_list, results) if ips - wildcard_ips}

tests/test_resolve.py

@@ -1,14 +1,24 @@
 """Tests for DNS resolution functionality.
 
 This module contains tests that verify the DNS resolution and domain
-registration checking works correctly for various domain inputs.
+registration checking works correctly for various domain inputs, including
+wildcard detection and wildcard-aware filtering.
 """
 
 import asyncio
 
-from subwiz.resolve import get_registered_domains
+from subwiz.resolve import (
+    NAME_SERVERS,
+    TIMEOUT,
+    TRIES,
+    _resolve_ips,
+    detect_wildcard,
+    get_registered_domains,
+)
 from subwiz.type import Domain
 
+import aiodns
+
 
 def test_():
     """Test that DNS resolution correctly identifies registered domains.
@@ -23,3 +33,47 @@ def test_():
         get_registered_domains(input_domains, resolution_concurrency=10)
     )
     assert registered_domains == {Domain("api.hadrian.io"), Domain("app.hadrian.io")}
+
+
+def test_detect_wildcard_non_wildcard():
+    """A domain without a wildcard record returns an empty catch-all set."""
+    wildcard_ips = asyncio.run(detect_wildcard("hadrian.io"))
+    assert wildcard_ips == set()
+
+
+def test_get_registered_domains_wildcard_keeps_real():
+    """Real subdomains resolving outside the wildcard set are kept.
+
+    Uses a TEST-NET address (RFC 5737, never routed to a real host) as the
+    wildcard set, so any real subdomain resolves outside it and is kept.
+    """
+    domain_strings = {"api.hadrian.io", "app.hadrian.io", "random_string.hadrian.io"}
+    input_domains = {Domain(dom) for dom in domain_strings}
+    registered_domains = asyncio.run(
+        get_registered_domains(
+            input_domains, resolution_concurrency=10, wildcard_ips={"192.0.2.1"}
+        )
+    )
+    assert registered_domains == {Domain("api.hadrian.io"), Domain("app.hadrian.io")}
+
+
+def test_get_registered_domains_wildcard_filters_catch_all():
+    """A subdomain resolving only to wildcard IPs is filtered out.
+
+    Learns api.hadrian.io's real IPs, then treats those exact IPs as the
+    wildcard set so the domain has no IP outside it and is dropped.
+    """
+    api = Domain("api.hadrian.io")
+
+    async def _run() -> set[Domain]:
+        resolver = aiodns.DNSResolver(
+            nameservers=NAME_SERVERS, timeout=TIMEOUT, tries=TRIES
+        )
+        semaphore = asyncio.Semaphore(1)
+        real_ips = await _resolve_ips(str(api), resolver, semaphore)
+        assert real_ips, "api.hadrian.io should resolve"
+        return await get_registered_domains(
+            {api}, resolution_concurrency=10, wildcard_ips=real_ips
+        )
+
+    assert asyncio.run(_run()) == set()