import asyncio
import aiodns
import csv
import random
import time

async def resolve_domain(resolver, domain, id):
    try:
        result = await resolver.query(domain, 'A')
        ip = result[0].host if result else 'No IP found'
        return id, domain, ip
    except Exception as e:
        return id, domain, f"Error: {e}"

async def resolve_domains_in_chunks(domain_id_pairs, resolvers, chunk_size=5000, sleep_time=5):
    results = []
    total_domains = len(domain_id_pairs)
    
    for i in range(0, total_domains, chunk_size):
        chunk = domain_id_pairs[i:i + chunk_size]
        tasks = [
            resolve_domain(random.choice(resolvers), domain, id) 
            for id, domain in chunk
        ]
        results.extend(await asyncio.gather(*tasks))

        # If there's more work to do, sleep between chunks
        if i + chunk_size < total_domains:
            print(f"Processed {i + chunk_size} out of {total_domains} domains. Sleeping for {sleep_time} seconds...")
            await asyncio.sleep(sleep_time)

    return results

def create_resolvers(dns_servers):
    resolvers = []
    for dns_server in dns_servers:
        resolver = aiodns.DNSResolver(nameservers=[dns_server])
        resolvers.append(resolver)
    return resolvers

def read_domains_from_csv(file_path):
    domain_id_pairs = []
    with open(file_path, mode='r', newline='', encoding='utf-8') as file:
        reader = csv.reader(file)
        for row in reader:
            if len(row) > 1:  # Ensure there are at least two columns
                id = row[0]   # Get the ID from the first column
                domain = row[1]  # Get the domain from the second column
                domain_id_pairs.append((id, domain))
    return domain_id_pairs

def save_results_to_file(results, output_file):
    with open(output_file, 'w') as f:
        for id, domain, ip in results:
            f.write(f"{id},{domain},{ip}\n")

def main():
    input_csv_file = 'tranco_V99PN_gTLDs_resolved_errors.csv'  # Replace with your CSV file path
    output_file = 'tranco_V99PN_gTLDs_resolved_2nd.csv'

    # Define your DNS servers
    dns_servers = [
        '8.8.8.8',    # Google Public DNS
        '8.8.4.4',    # Google Secondary DNS
        #'1.1.1.1',    # Cloudflare Public DNS
        #'1.1.0.1',    # Cloudflare Secondary DNS
        '9.9.9.10',   # Quad9 No Malware Blocking
        '149.112.112.10' # Quad9 No Malware Blocking Secondary
        #'208.67.222.222', # OpenDNS Primary
        #'208.67.220.220'  # OpenDNS Secondary
        # Add more DNS servers if needed
    ]

    # Create resolvers for each DNS server
    resolvers = create_resolvers(dns_servers)

    # Read domains and IDs from the CSV file
    domain_id_pairs = read_domains_from_csv(input_csv_file)
    
    if not domain_id_pairs:
        print("No domains found in the CSV file.")
        return
    
    loop = asyncio.get_event_loop()
    results = loop.run_until_complete(resolve_domains_in_chunks(domain_id_pairs, resolvers, chunk_size=5000, sleep_time=5))
    
    save_results_to_file(results, output_file)
    print(f"DNS resolution complete. Results saved to {output_file}")

if __name__ == "__main__":
    main()