From 46d0d464ed65b2ba2cde71aa7d9d6c3593b2661db30f94c6c8852d6bcbbc6d76 Mon Sep 17 00:00:00 2001 From: Discsearcher Date: Tue, 11 Aug 2026 23:46:12 -0400 Subject: [PATCH] more updates like timeouts and threading. --- dns_compare.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/dns_compare.py b/dns_compare.py index 2ae60fa..1136c25 100755 --- a/dns_compare.py +++ b/dns_compare.py @@ -33,6 +33,8 @@ parser.add_argument('--nameservers', nargs='+', help='IPv4 or IPv6 address of na parser.add_argument('--ns-builtin', nargs='+', help=f'Built-in public DNS servers. Options: {", ".join(builtin_nameservers.keys())}, all') parser.add_argument('--records', nargs='*', help='Hostname(s) to query. If not provided with --domain, queries the domain itself.') parser.add_argument('--record-type', default='A', help='DNS record type to query. Default: A') +parser.add_argument('--threads', type=int, default=None, help='Number of threads to use for queries. Default: CPU count - 1') +parser.add_argument('--timeout', type=float, default=5.0, help='DNS query timeout in seconds. Default: 5.0') parser.add_argument('--list-record-types', action='store_true', help='Print list of supported record types and exit.') args = parser.parse_args() @@ -87,18 +89,18 @@ except rdatatype.UnknownRdatatype: exit(1) # Function to perform a single DNS query -def perform_query(full_domain, ns, record_type): +def perform_query(full_domain, ns, record_type, timeout): """Perform a DNS query and return results""" try: qname = name.from_text(full_domain) q = message.make_query(qname, record_type) # Try UDP first - r = query.udp(q, ns) + r = query.udp(q, ns, timeout=timeout) # If truncated, retry with TCP to get all results if r.flags & flags.TC: - r = query.tcp(q, ns) + r = query.tcp(q, ns, timeout=timeout) results = [] answers = r.answer @@ -140,13 +142,13 @@ for full_domain in full_domains: for ns in all_nameservers: query_tasks.append((full_domain, ns, record_type)) -# Determine thread count: CPU count - 1, minimum 1 -thread_count = 1 if 1 else os.cpu_count() - 1 +# Determine thread count: use provided value or default to CPU count - 1, minimum 1 +thread_count = args.threads if args.threads else max(1, os.cpu_count() - 1) # Execute queries in parallel results_map = {} # (full_domain, ns) -> resolved_ip with ThreadPoolExecutor(max_workers=thread_count) as executor: - futures = [executor.submit(perform_query, full_domain, ns, record_type) + futures = [executor.submit(perform_query, full_domain, ns, record_type, args.timeout) for full_domain, ns, _ in query_tasks] for future in futures: