From adcf82f14c293fa447ccedc500f0ee6ddae74807 Mon Sep 17 00:00:00 2001 From: Parth Shukla Date: Mon, 7 Sep 2026 12:41:00 +0530 Subject: [PATCH 1/2] fix: detect wildcard DNS up front and filter false positives Domains with a wildcard record (*.example.com) resolve every predicted subdomain to the same catch-all IPs, producing huge false-positive counts that compound through the recursive inference loop. This detects wildcards once per apex, up front (before the expensive transformer inference), by probing random subdomains. When a wildcard is present, resolution keeps only predictions that resolve to at least one IP outside the catch-all set; predictions resolving solely to the wildcard IPs are dropped as noise. A new --wildcard {filter,skip} flag (default: filter) lets the user skip wildcard apexes entirely instead. - resolve.py: add detect_wildcard(), _resolve_ips() helper; get_registered_domains() gains an optional wildcard_ips filter. - main.py: detect wildcard per apex in run(), warn, and thread wildcard_ips through to resolution; add wildcard param + validation. - cli.py: add --wildcard flag. - tests: cover wildcard detection and wildcard-aware filtering. - README: regenerate CLI help block (adds --wildcard; reflects current argparse). - model.py: black formatting only (no logic change). Co-Authored-By: Claude Opus 4.8 --- README.md | 9 +++- subwiz/cli.py | 8 +++ subwiz/main.py | 42 ++++++++++++++-- subwiz/model.py | 1 - subwiz/resolve.py | 111 +++++++++++++++++++++++++++++++++++++++--- tests/test_resolve.py | 58 +++++++++++++++++++++- 6 files changed, 213 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index 7190f02..6442214 100644 --- a/README.md +++ b/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) ``` diff --git a/subwiz/cli.py b/subwiz/cli.py index 78f517e..dfb7b94 100644 --- a/subwiz/cli.py +++ b/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", diff --git a/subwiz/main.py b/subwiz/main.py index 0910589..8cebd64 100644 --- a/subwiz/main.py +++ b/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,28 @@ def run( found_domains = set() for apex in sorted(domain_groups): + # Detect wildcards once, up front, before the expensive inference. A + # wildcard makes plain DNS resolution meaningless (every subdomain + # resolves), so we either skip the apex or filter to subdomains that + # resolve outside the catch-all IP set. + 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 +369,7 @@ def run( no_resolve=no_resolve, resolution_concurrency=resolution_concurrency, quiet=quiet, + wildcard_ips=wildcard_ips, ) return sorted(found_domains) diff --git a/subwiz/model.py b/subwiz/model.py index 5295c9b..36423e9 100644 --- a/subwiz/model.py +++ b/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])?))*(? 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] + ) + + # A wildcard exists only if every random probe resolved. If any random + # label fails to resolve, there is no catch-all record. + 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} diff --git a/tests/test_resolve.py b/tests/test_resolve.py index deb5e14..82f75f8 100644 --- a/tests/test_resolve.py +++ b/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() From 5cb4538c2e35bd088fc522dbcea1aaaab7fa0a4f Mon Sep 17 00:00:00 2001 From: Parth Shukla Date: Tue, 15 Sep 2026 13:46:37 +0530 Subject: [PATCH 2/2] refactor: remove verbose inline comments per review Addresses review feedback on #38. Co-Authored-By: Claude Opus 4.8 --- subwiz/main.py | 4 ---- subwiz/resolve.py | 5 ----- 2 files changed, 9 deletions(-) diff --git a/subwiz/main.py b/subwiz/main.py index 8cebd64..f67cb67 100644 --- a/subwiz/main.py +++ b/subwiz/main.py @@ -334,10 +334,6 @@ def run( found_domains = set() for apex in sorted(domain_groups): - # Detect wildcards once, up front, before the expensive inference. A - # wildcard makes plain DNS resolution meaningless (every subdomain - # resolves), so we either skip the apex or filter to subdomains that - # resolve outside the catch-all IP set. wildcard_ips = None if not no_resolve: wildcard_ips = asyncio.run(detect_wildcard(apex, resolution_concurrency)) diff --git a/subwiz/resolve.py b/subwiz/resolve.py index 0bbd680..76380ae 100644 --- a/subwiz/resolve.py +++ b/subwiz/resolve.py @@ -26,9 +26,6 @@ TIMEOUT = 3 TRIES = 1 -# Number of random subdomains probed to detect a wildcard record, and the -# length of each random label. Labels are long enough that a real subdomain -# collision is effectively impossible. WILDCARD_PROBE_COUNT = 3 WILDCARD_LABEL_LENGTH = 20 @@ -93,8 +90,6 @@ async def detect_wildcard( *[_resolve_ips(probe, resolver, semaphore) for probe in probes] ) - # A wildcard exists only if every random probe resolved. If any random - # label fails to resolve, there is no catch-all record. if not all(probe_ips): return set()