diff --git a/dns_compare.py b/dns_compare.py index 1136c25..a003770 100755 --- a/dns_compare.py +++ b/dns_compare.py @@ -11,6 +11,7 @@ from dns import query from dns import rdatatype from dns import name from dns import flags +from dns import exception # Get list of supported record types supported_types = sorted([rdatatype.to_text(rdtype) for rdtype in rdatatype.RdataType]) @@ -110,8 +111,11 @@ def perform_query(full_domain, ns, record_type, timeout): results.append(str(item)) return (full_domain, ns, results) + except (exception.Timeout, TimeoutError): + return (full_domain, ns, "TIMED OUT") except Exception as e: - return (full_domain, ns, None) + error_msg = str(e) if str(e) else type(e).__name__ + return (full_domain, ns, f"ERROR: {error_msg}") # Start timer start_time = time.time() @@ -153,7 +157,7 @@ with ThreadPoolExecutor(max_workers=thread_count) as executor: for future in futures: full_domain, ns, resolved_ips = future.result() - if resolved_ips: + if resolved_ips is not None: results_map[(full_domain, ns)] = resolved_ips # Build the output structure @@ -179,16 +183,26 @@ for full_domain in full_domains: for ns in all_nameservers: provider = nameserver_map.get(ns) if (full_domain, ns) in results_map: - results = sorted(results_map[(full_domain, ns)]) - # Store as single value if only one result, otherwise as array - output[full_domain]["records"][provider][ns] = results[0] if len(results) == 1 else results + results = results_map[(full_domain, ns)] + # Handle error/timeout cases (strings starting with specific markers) + if isinstance(results, str) and (results == "TIMED OUT" or results.startswith("ERROR:")): + output[full_domain]["records"][provider][ns] = results + else: + results = sorted(results) + # Store as single value if only one result, otherwise as array + output[full_domain]["records"][provider][ns] = results[0] if len(results) == 1 else results else: # Flat structure for manual nameservers only for ns in all_nameservers: if (full_domain, ns) in results_map: - results = sorted(results_map[(full_domain, ns)]) - # Store as single value if only one result, otherwise as array - output[full_domain]["records"][ns] = results[0] if len(results) == 1 else results + results = results_map[(full_domain, ns)] + # Handle error/timeout cases (strings starting with specific markers) + if isinstance(results, str) and (results == "TIMED OUT" or results.startswith("ERROR:")): + output[full_domain]["records"][ns] = results + else: + results = sorted(results) + # Store as single value if only one result, otherwise as array + output[full_domain]["records"][ns] = results[0] if len(results) == 1 else results # Calculate and add timing information end_time = time.time()