<commit_before>from base64 import b64encode
import dns.rdatatype
import dns.message
import dns.resolver
import select
import socket
import threading
import time


def get_ips(host, nameserver=None, record="A"):
    nameservers = []
    if nameserver is not None:
        nameservers = [nameserver]
    return lookup_domain(host, nameservers=nameservers, rtype=record)


def lookup_domain(domain, nameservers=[], rtype="A", timeout=10):
    """Wrapper for DNSQuery method"""
    dns_exp = DNSQuery(domains=[domain], nameservers=nameservers, rtype=rtype,
                       timeout=timeout)
    return dns_exp.lookup_domain(domain)


def lookup_domains(domains, nameservers=[], rtype="A", timeout=10):
    dns_exp = DNSQuery(domains=domains, nameservers=nameservers, rtype=rtype,
                       timeout=timeout)
    return dns_exp.lookup_domains()


def send_chaos_queries():
    dns_exp = DNSQuery()
    return dns_exp.send_chaos_queries()


class DNSQuery():
    """Class to store state for all of the DNS queries"""

    def __init__(self, domains=[], nameservers=[], rtype="A", timeout=10,
                 max_threads=100):
        """Constructor for the DNS query class

        Params:
        nameserver- the nameserver to use, defaults to the local resolver
        rtype- the record type to lookup (as text), by default A
        timeout- how long to wait for a response, by default 10 seconds

        """
        self.domains = domains
        self.rtype = rtype
        self.timeout = timeout
        self.max_threads = max_threads
        if nameservers == []:
            nameservers = dns.resolver.Resolver().nameservers
        self.nameservers = nameservers
        self.results = {}
        self.threads = []

    def send_chaos_queries(self):
        """Send chaos queries to identify the DNS server and its manufacturer

        Note: we send 2 queries for BIND stuff per RFC 4892 and 1
        query per RFC 6304

        Note: we are not waiting on a second response because we
        shouldn't be getting injected packets here

        """
        names = ["HOSTNAME.BIND", "VERSION.BIND", "ID.SERVER"]
        self.results = {'exp-name': "chaos-queries"}
        for name in names:
            self.results[name] = {}
            for nameserver in self.nameservers:
                sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                sock.settimeout(self.timeout)
                query = dns.message.make_query(name,
                                               dns.rdatatype.from_text("TXT"),
                                               dns.rdataclass.from_text("CH"))
                sock.sendto(query.to_wire(), (nameserver, 53))
                reads, _, _ = select.select([sock], [], [], self.timeout)
                if reads == []:
                    self.results[name][nameserver] = None
                else:
                    response = reads[0].recvfrom(4096)[0]
                    self.results[name][nameserver] = b64encode(response)
        return self.results

    def lookup_domains(self):
        """More complex DNS primitive that lookups domains concurrently

        Note: if you want to lookup multiple domains, you should use
        this function
        """
        thread_error = False
        thread_wait_timeout = 200
        for nameserver in self.nameservers:
            for domain in self.domains:
                while threading.active_count() > self.max_threads:
                    time.sleep(1)
                    wait_time += 1
                    if wait_time > thread_wait_timeout:
                        thread_error = True
                        break

                if thread_error:
                    self.results["error"] = "Threads took too long to finish."
                    break

                thread = threading.Thread(target=self.lookup_domain,
                                          args=(domain, nameserver))
                thread.setDaemon(1)
                thread.start()
                self.threads.append(thread)<commit_msg>fixed a small bug in DNS primitive<commit_after>from base64 import b64encode
import dns.rdatatype
import dns.message
import dns.resolver
import select
import socket
import threading
import time


def get_ips(host, nameserver=None, record="A"):
    nameservers = []
    if nameserver is not None:
        nameservers = [nameserver]
    return lookup_domain(host, nameservers=nameservers, rtype=record)


def lookup_domain(domain, nameservers=[], rtype="A", timeout=10):
    """Wrapper for DNSQuery method"""
    dns_exp = DNSQuery(domains=[domain], nameservers=nameservers, rtype=rtype,
                       timeout=timeout)
    return dns_exp.lookup_domain(domain)


def lookup_domains(domains, nameservers=[], rtype="A", timeout=10):
    dns_exp = DNSQuery(domains=domains, nameservers=nameservers, rtype=rtype,
                       timeout=timeout)
    return dns_exp.lookup_domains()


def send_chaos_queries():
    dns_exp = DNSQuery()
    return dns_exp.send_chaos_queries()


class DNSQuery():
    """Class to store state for all of the DNS queries"""

    def __init__(self, domains=[], nameservers=[], rtype="A", timeout=10,
                 max_threads=100):
        """Constructor for the DNS query class

        Params:
        nameserver- the nameserver to use, defaults to the local resolver
        rtype- the record type to lookup (as text), by default A
        timeout- how long to wait for a response, by default 10 seconds

        """
        self.domains = domains
        self.rtype = rtype
        self.timeout = timeout
        self.max_threads = max_threads
        if nameservers == []:
            nameservers = dns.resolver.Resolver().nameservers
        self.nameservers = nameservers
        self.results = {}
        self.threads = []

    def send_chaos_queries(self):
        """Send chaos queries to identify the DNS server and its manufacturer

        Note: we send 2 queries for BIND stuff per RFC 4892 and 1
        query per RFC 6304

        Note: we are not waiting on a second response because we
        shouldn't be getting injected packets here

        """
        names = ["HOSTNAME.BIND", "VERSION.BIND", "ID.SERVER"]
        self.results = {'exp-name': "chaos-queries"}
        for name in names:
            self.results[name] = {}
            for nameserver in self.nameservers:
                sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                sock.settimeout(self.timeout)
                query = dns.message.make_query(name,
                                               dns.rdatatype.from_text("TXT"),
                                               dns.rdataclass.from_text("CH"))
                sock.sendto(query.to_wire(), (nameserver, 53))
                reads, _, _ = select.select([sock], [], [], self.timeout)
                if reads == []:
                    self.results[name][nameserver] = None
                else:
                    response = reads[0].recvfrom(4096)[0]
                    self.results[name][nameserver] = b64encode(response)
        return self.results

    def lookup_domains(self):
        """More complex DNS primitive that lookups domains concurrently

        Note: if you want to lookup multiple domains, you should use
        this function
        """
        thread_error = False
        thread_wait_timeout = 200
        for nameserver in self.nameservers:
            for domain in self.domains:
                wait_time = 0
                while threading.active_count() > self.max_threads:
                    time.sleep(1)
                    wait_time += 1
                    if wait_time > thread_wait_timeout:
                        thread_error = True
                        break

                if thread_error:
                    self.results["error"] = "Threads took too long to finish."
                    break

                thread = threading.Thread(target=self.lookup_domain,
                                          args=(domain, nameserver))
                thread.setDaemon(1)
                thread.start()
                self.threads.append(thread)
{'fraction_non_alphanumeric': 0.08883279240293304, 'fraction_numerical': 0.010217574227671595, 'mean_word_length': 1.9111266620013996, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1528101', 'n_tokens_mistral': 2375, 'n_tokens_neox': 2158, 'n_words': 603}