Example #1
0
    def _init_udp(self):
        """
        Create a new Datagram instance and listen on a socket.
        """
        self._udp = Datagram(engine=self.engine)
        self._udp.on_read = self.receive_message

        start = port = random.randrange(10005, 65535)
        while True:
            try:
                self._udp.listen(('', port))
                break
            except Exception:
                port += 1
                if port > 65535:
                    port = 10000
                if port == start:
                    raise Exception("Can't listen on any port.")
Example #2
0
File: dns.py Project: ixokai/pants
    def _init_udp(self):
        """
        Create a new Datagram instance and listen on a socket.
        """
        self._udp = Datagram()
        self._udp.on_read = self.receive_message

        start = port = random.randrange(10005, 65535)
        while True:
            try:
                self._udp.listen(("", port))
                break
            except Exception:
                port += 1
                if port > 65535:
                    port = 10000
                if port == start:
                    raise Exception("Can't listen on any port.")
Example #3
0
File: dns.py Project: ixokai/pants
class Resolver(object):
    """
    The Resolver class generates DNS messages, sends them to remote servers,
    and processes any responses. The bulk of the heavy lifting is done in
    DNSMessage and the RDATA handling functions, however.

    =========  ============
    Argument   Description
    =========  ============
    servers    *Optional.* A list of DNS servers to query. If a list isn't provided, Pants will attempt to retrieve a list of servers from the OS, falling back to a list of default servers if none are available.
    =========  ============
    """

    def __init__(self, servers=None):
        self.servers = servers or list_dns_servers()

        # Internal State
        self._messages = {}
        self._cache = {}
        self._queries = {}
        self._tcp = {}
        self._udp = None
        self._last_id = -1

    def _safely_call(self, callback, *args, **kwargs):
        try:
            callback(*args, **kwargs)
        except Exception:
            log.exception("Error calling callback for DNS result.")

    def _error(self, message, err=DNS_TIMEOUT):
        if not message in self._messages:
            return

        if message in self._tcp:
            try:
                self._tcp[message].close()
            except Exception:
                pass
            del self._tcp[message]

        callback, message, df_timeout, media, data = self._messages[message]
        del self._messages[message.id]

        try:
            df_timeout.cancel()
        except Exception:
            pass

        if err == DNS_TIMEOUT and data:
            self._safely_call(callback, DNS_OK, data)
        else:
            self._safely_call(callback, err, None)

    def _init_udp(self):
        """
        Create a new Datagram instance and listen on a socket.
        """
        self._udp = Datagram()
        self._udp.on_read = self.receive_message

        start = port = random.randrange(10005, 65535)
        while True:
            try:
                self._udp.listen(("", port))
                break
            except Exception:
                port += 1
                if port > 65535:
                    port = 10000
                if port == start:
                    raise Exception("Can't listen on any port.")

    def send_message(self, message, callback=None, timeout=10, media=None):
        """
        Send an instance of DNSMessage to a DNS server, and call the provided
        callback when a response is received, or if the action times out.

        =========  ========  ============
        Argument   Default   Description
        =========  ========  ============
        message              The :class:`DNSMessage` instance to send to the server.
        callback   None      *Optional.* The function to call once the response has been received or the attempt has timed out.
        timeout    10        *Optional.* How long, in seconds, to wait before timing out.
        media      None      *Optional.* Whether to use UDP or TCP. UDP is used by default.
        =========  ========  ============
        """
        while message.id is None or message.id in self._messages:
            self._last_id += 1
            if self._last_id > 65535:
                self._last_id = 0
            message.id = self._last_id

        # Timeout in timeout seconds.
        df_timeout = pants.engine.defer(timeout, self._error, message.id)

        # Send the Message
        msg = str(message)
        if media is None:
            media = "udp"
            # if len(msg) > 512:
            #    media = 'tcp'
            # else:
            #    media = 'udp'

        # Store Info
        self._messages[message.id] = callback, message, df_timeout, media, None

        if media == "udp":
            if self._udp is None:
                self._init_udp()
            try:
                self._udp.write(msg, (self.servers[0], DNS_PORT))
            except Exception:
                # Pants gummed up. Try again.
                self._next_server(message.id)

            pants.engine.defer(0.5, self._next_server, message.id)

        else:
            tcp = self._tcp[message.id] = _DNSStream(self, message.id)
            tcp.connect((self.servers[0], DNS_PORT))

    def _next_server(self, id):
        if not id in self._messages or id in self._tcp:
            return

        # Cycle the list.
        self.servers.append(self.servers.pop(0))

        msg = str(self._messages[id][1])
        try:
            self._udp.write(msg, (self.servers[0], DNS_PORT))
        except Exception:
            try:
                self._udp.close()
            except Exception:
                pass
            del self._udp
            self._init_udp()
            self._udp.write(msg, (self.servers[0], DNS_PORT))

    def receive_message(self, data):
        if not isinstance(data, DNSMessage):
            try:
                data = DNSMessage.from_string(data)
            except TooShortError:
                if len(data) < 2:
                    return

                id = struct.unpack("!H", data[:2])
                if not id in self._messages:
                    return

                self._error(id, err=DNS_BADRESPONSE)
                return

        if not data.id in self._messages:
            return

        callback, message, df_timeout, media, _ = self._messages[data.id]

        # if data.tc and media == 'udp':
        #    self._messages[data.id] = callback, message, df_timeout, 'tcp', data
        #    tcp = self._tcp[data.id] = _DNSStream(self, message.id)
        #    tcp.connect((self.servers[0], DNS_PORT))
        #    return

        if not data.server:
            if self._udp and isinstance(self._udp.remote_addr, tuple):
                data.server = "%s:%d" % self._udp.remote_addr
            else:
                data.server = "%s:%d" % (self.servers[0], DNS_PORT)

        try:
            df_timeout.cancel()
        except Exception:
            pass

        del self._messages[data.id]
        self._safely_call(callback, DNS_OK, data)

    def query(self, name, qtype=A, qclass=IN, callback=None, timeout=10, allow_cache=True, allow_hosts=True):
        """
        Make a DNS request of the given QTYPE for the given name.

        ============  ========  ============
        Argument      Default   Description
        ============  ========  ============
        name                    The name to query.
        qtype         A         *Optional.* The QTYPE to query.
        qclass        IN        *Optional.* The QCLASS to query.
        callback      None      *Optional.* The function to call when a response for the query has been received, or when the request has timed out.
        timeout       10        *Optional.* The time, in seconds, to wait before timing out.
        allow_cache   True      *Optional.* Whether or not to use the cache. If you expect to be performing thousands of requests, you may want to disable the cache to avoid excess memory usage.
        allow_hosts   True      *Optional.* Whether or not to use any records gathered from the OS hosts file.
        ============  ========  ============
        """
        if allow_hosts:
            if host_time + 30 < time.time():
                load_hosts()

            cname = None
            if name in self._cache and CNAME in self._cache[name]:
                cname = self._cache[name][CNAME]

            if qtype == A and name in hosts[A]:
                self._safely_call(callback, DNS_OK, cname, None, (hosts[A][name],))

            elif qtype == AAAA and name in hosts[AAAA]:
                self._safely_call(callback, DNS_OK, cname, None, (hosts[AAAA][name],))

        if allow_cache and name in self._cache and (qtype, qclass) in self._cache[name]:
            cname = None
            if CNAME in self._cache[name]:
                cname = self._cache[name][CNAME]
            death, ttl, rdata = self._cache[name][(qtype, qclass)]

            if death < time.time():
                # Clear out the old record.
                del self._cache[name][(qtype, qclass)]

            else:
                if callback:
                    self._safely_call(callback, DNS_OK, cname, ttl, rdata)
                return

        # Build a message and add our question.
        m = DNSMessage()
        m.questions.append((name, qtype, qclass))

        # Make the function for handling our response.
        def handle_response(status, data):
            cname = None
            # TTL is 30 by default, so answers with no records we want will be
            # repeated, but not too often.
            ttl = 30

            if not data:
                self._safely_call(callback, status, None, None, None)
                return

            rdata = []
            for (aname, atype, aclass, attl, ardata) in data.answers:
                if atype == CNAME:
                    cname = ardata[0]

                if atype == qtype and aclass == qclass:
                    ttl = attl
                    if len(ardata) == 1:
                        rdata.append(ardata[0])
                    else:
                        rdata.append(ardata)
            rdata = tuple(rdata)

            if allow_cache:
                if not name in self._cache:
                    self._cache[name] = {}
                    if cname:
                        self._cache[name][CNAME] = cname
                    self._cache[name][(qtype, qclass)] = time.time() + ttl, ttl, rdata

            if data.rcode != DNS_OK:
                status = data.rcode

            self._safely_call(callback, status, cname, ttl, rdata)

        # Send it, so we get an ID.
        self.send_message(m, handle_response)
Example #4
0
class Resolver(object):
    """
    The Resolver class generates DNS messages, sends them to remote servers,
    and processes any responses. The bulk of the heavy lifting is done in
    DNSMessage and the RDATA handling functions, however.

    =========  ============
    Argument   Description
    =========  ============
    servers    *Optional.* A list of DNS servers to query. If a list isn't provided, Pants will attempt to retrieve a list of servers from the OS, falling back to a list of default servers if none are available.
    engine     *Optional.* The :class:`pants.engine.Engine` instance to use.
    =========  ============
    """
    def __init__(self, servers=None, engine=None):
        self.servers = servers or list_dns_servers()
        self.engine = engine or Engine.instance()

        # Internal State
        self._messages = {}
        self._cache = {}
        self._queries = {}
        self._tcp = {}
        self._udp = None
        self._last_id = -1

    def _safely_call(self, callback, *args, **kwargs):
        try:
            callback(*args, **kwargs)
        except Exception:
            log.exception('Error calling callback for DNS result.')

    def _error(self, message, err=DNS_TIMEOUT):
        if not message in self._messages:
            return

        if message in self._tcp:
            try:
                self._tcp[message].close()
            except Exception:
                pass
            del self._tcp[message]

        callback, message, df_timeout, media, data = self._messages[message]
        del self._messages[message.id]

        try:
            df_timeout.cancel()
        except Exception:
            pass

        if err == DNS_TIMEOUT and data:
            self._safely_call(callback, DNS_OK, data)
        else:
            self._safely_call(callback, err, None)

    def _init_udp(self):
        """
        Create a new Datagram instance and listen on a socket.
        """
        self._udp = Datagram(engine=self.engine)
        self._udp.on_read = self.receive_message

        start = port = random.randrange(10005, 65535)
        while True:
            try:
                self._udp.listen(('', port))
                break
            except Exception:
                port += 1
                if port > 65535:
                    port = 10000
                if port == start:
                    raise Exception("Can't listen on any port.")

    def send_message(self, message, callback=None, timeout=10, media=None):
        """
        Send an instance of DNSMessage to a DNS server, and call the provided
        callback when a response is received, or if the action times out.

        =========  ========  ============
        Argument   Default   Description
        =========  ========  ============
        message              The :class:`DNSMessage` instance to send to the server.
        callback   None      *Optional.* The function to call once the response has been received or the attempt has timed out.
        timeout    10        *Optional.* How long, in seconds, to wait before timing out.
        media      None      *Optional.* Whether to use UDP or TCP. UDP is used by default.
        =========  ========  ============
        """
        while message.id is None or message.id in self._messages:
            self._last_id += 1
            if self._last_id > 65535:
                self._last_id = 0
            message.id = self._last_id

        # Timeout in timeout seconds.
        df_timeout = self.engine.defer(timeout, self._error, message.id)

        # Send the Message
        msg = str(message)
        if media is None:
            media = 'udp'
            #if len(msg) > 512:
            #    media = 'tcp'
            #else:
            #    media = 'udp'

        # Store Info
        self._messages[message.id] = callback, message, df_timeout, media, None

        if media == 'udp':
            if self._udp is None:
                self._init_udp()
            try:
                self._udp.write(msg, (self.servers[0], DNS_PORT))
            except Exception:
                # Pants gummed up. Try again.
                self._next_server(message.id)

            self.engine.defer(0.5, self._next_server, message.id)

        else:
            tcp = self._tcp[message.id] = _DNSStream(self, message.id)
            tcp.connect((self.servers[0], DNS_PORT))

    def _next_server(self, id):
        if not id in self._messages or id in self._tcp:
            return

        # Cycle the list.
        self.servers.append(self.servers.pop(0))

        msg = str(self._messages[id][1])
        try:
            self._udp.write(msg, (self.servers[0], DNS_PORT))
        except Exception:
            try:
                self._udp.close()
            except Exception:
                pass
            del self._udp
            self._init_udp()
            self._udp.write(msg, (self.servers[0], DNS_PORT))

    def receive_message(self, data):
        if not isinstance(data, DNSMessage):
            try:
                data = DNSMessage.from_string(data)
            except TooShortError:
                if len(data) < 2:
                    return

                id = struct.unpack("!H", data[:2])
                if not id in self._messages:
                    return

                self._error(id, err=DNS_FORMATERROR)
                return

        if not data.id in self._messages:
            return

        callback, message, df_timeout, media, _ = self._messages[data.id]

        #if data.tc and media == 'udp':
        #    self._messages[data.id] = callback, message, df_timeout, 'tcp', data
        #    tcp = self._tcp[data.id] = _DNSStream(self, message.id)
        #    tcp.connect((self.servers[0], DNS_PORT))
        #    return

        if not data.server:
            if self._udp and isinstance(self._udp.remote_address, tuple):
                data.server = '%s:%d' % self._udp.remote_address
            else:
                data.server = '%s:%d' % (self.servers[0], DNS_PORT)

        try:
            df_timeout.cancel()
        except Exception:
            pass

        del self._messages[data.id]
        self._safely_call(callback, DNS_OK, data)

    def query(self,
              name,
              qtype=A,
              qclass=IN,
              callback=None,
              timeout=10,
              allow_cache=True,
              allow_hosts=True):
        """
        Make a DNS request of the given QTYPE for the given name.

        ============  ========  ============
        Argument      Default   Description
        ============  ========  ============
        name                    The name to query.
        qtype         A         *Optional.* The QTYPE to query.
        qclass        IN        *Optional.* The QCLASS to query.
        callback      None      *Optional.* The function to call when a response for the query has been received, or when the request has timed out.
        timeout       10        *Optional.* The time, in seconds, to wait before timing out.
        allow_cache   True      *Optional.* Whether or not to use the cache. If you expect to be performing thousands of requests, you may want to disable the cache to avoid excess memory usage.
        allow_hosts   True      *Optional.* Whether or not to use any records gathered from the OS hosts file.
        ============  ========  ============
        """
        if not isinstance(qtype, (list, tuple)):
            qtype = (qtype, )

        if allow_hosts:
            if host_time + 30 < time.time():
                load_hosts()

            cname = None
            if name in self._cache and CNAME in self._cache[name]:
                cname = self._cache[name][CNAME]

            result = []

            if AAAA in qtype and name in hosts[AAAA]:
                result.append(hosts[AAAA][name])

            if A in qtype and name in hosts[A]:
                result.append(hosts[A][name])

            if result:
                if callback:
                    self._safely_call(callback, DNS_OK, cname, None,
                                      tuple(result))
                return

        if allow_cache and name in self._cache:
            cname = self._cache[name].get(CNAME, None)

            tm = time.time()
            result = []
            min_ttl = sys.maxint

            for t in qtype:
                death, ttl, rdata = self._cache[name][(t, qclass)]
                if death < tm:
                    del self._cache[name][(t, qclass)]
                    continue

                min_ttl = min(ttl, min_ttl)
                if rdata:
                    result.extend(rdata)

            if callback:
                self._safely_call(callback, DNS_OK, cname, min_ttl,
                                  tuple(result))
            return

        # Build a message and add our question.
        m = DNSMessage()

        m.questions.append((name, qtype[0], qclass))

        # Make the function for handling our response.
        def handle_response(status, data):
            cname = None
            # TTL is 30 by default, so answers with no records we want will be
            # repeated, but not too often.
            ttl = sys.maxint

            if not data:
                self._safely_call(callback, status, None, None, None)
                return

            rdata = {}
            final_rdata = []
            for (aname, atype, aclass, attl, ardata) in data.answers:
                if atype == CNAME:
                    cname = ardata[0]

                if atype in qtype and aclass == qclass:
                    ttl = min(attl, ttl)
                    if len(ardata) == 1:
                        rdata.setdefault(atype, []).append(ardata[0])
                        final_rdata.append(ardata[0])
                    else:
                        rdata.setdefault(atype, []).append(ardata)
                        final_rdata.append(ardata)
            final_rdata = tuple(final_rdata)
            ttl = min(30, ttl)

            if allow_cache:
                if not name in self._cache:
                    self._cache[name] = {}
                    if cname:
                        self._cache[name][CNAME] = cname
                    for t in qtype:
                        self._cache[name][(
                            t, qclass)] = time.time() + ttl, ttl, rdata.get(
                                t, [])

            if data.rcode != DNS_OK:
                status = data.rcode

            self._safely_call(callback, status, cname, ttl, final_rdata)

        # Send it, so we get an ID.
        self.send_message(m, handle_response)