more updates like timeouts and threading.

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
+# Determine thread count: use provided value or default to CPU count - 1, minimum 1
-thread_count = 1 if 1 else os.cpu_count() - 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: