# -*- Mode: Python; tab-width: 4 -*-
# Author: Sam Rushing <rushing@nightmare.com>

# Fast, low-overhead asynchronous name resolver.  uses 'pre-cooked'
# DNS requests, unpacks only as much as it needs of the reply.

# see rfc1035 for details

import string
import asyncore
import socket
import sys
import time
from counter import counter

RCS_ID = '$Id$'
VERSION = '0.0'

# header
#                                    1  1  1  1  1  1
#      0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#    |                      ID                       |
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#    |QR|   Opcode  |AA|TC|RD|RA|   Z    |   RCODE   |
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#    |                    QDCOUNT                    |
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#    |                    ANCOUNT                    |
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#    |                    NSCOUNT                    |
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#    |                    ARCOUNT                    |
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+


# question
#                                    1  1  1  1  1  1
#      0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#    |                                               |
#    /                     QNAME                     /
#    /                                               /
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#    |                     QTYPE                     |
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#    |                     QCLASS                    |
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+

# build a DNS address request, _quickly_
def fast_address_request(host, id=0):
    return (
        '%c%c' % (chr((id >> 8) & 0xff), chr(id & 0xff)) +
        '\001\000\000\001\000\000\000\000\000\000%s\000\000\001\000\001' % (
            string.join(
                map(
                    lambda part: '%c%s' % (chr(len(part)), part),
                    string.split(host, '.')
                ), ''
            )
        )
    )


def fast_ptr_request(host, id=0):
    return (
        '%c%c' % (chr((id >> 8) & 0xff), chr(id & 0xff)) +
        '\001\000\000\001\000\000\000\000\000\000%s\000\000\014\000\001' % (
            string.join(
                map(
                    lambda part: '%c%s' % (chr(len(part)), part),
                    string.split(host, '.')
                ), ''
            )
        )
    )


def unpack_name(r, pos):
    n = []
    while 1:
        ll = ord(r[pos])
        if (ll & 0xc0):
            # compression
            pos = (ll & 0x3f << 8) + (ord(r[pos + 1]))
        elif ll == 0:
            break
        else:
            pos = pos + 1
            n.append(r[pos:pos + ll])
            pos = pos + ll
    return string.join(n, '.')


def skip_name(r, pos):
    while 1:
        ll = ord(r[pos])
        if (ll & 0xc0):
            # compression
            return pos + 2
        elif ll == 0:
            pos = pos + 1
            break
        else:
            pos = pos + ll + 1
    return pos


def unpack_ttl(r, pos):
    return reduce(
        lambda x, y: (x << 8) | y,
        map(ord, r[pos:pos + 4])
    )


# resource record
#                                    1  1  1  1  1  1
#      0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#    |                                               |
#    /                                               /
#    /                      NAME                     /
#    |                                               |
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#    |                      TYPE                     |
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#    |                     CLASS                     |
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#    |                      TTL                      |
#    |                                               |
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#    |                   RDLENGTH                    |
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--|
#    /                     RDATA                     /
#    /                                               /
#    +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+

def unpack_address_reply(r):
    ancount = (ord(r[6]) << 8) + (ord(r[7]))
    # skip question, first name starts at 12,
    # this is followed by QTYPE and QCLASS
    pos = skip_name(r, 12) + 4
    if ancount:
        # we are looking very specifically for
        # an answer with TYPE=A, CLASS=IN (\000\001\000\001)
        for an in range(ancount):
            pos = skip_name(r, pos)
            if r[pos:pos + 4] == '\000\001\000\001':
                return (
                    unpack_ttl(r, pos + 4),
                    '%d.%d.%d.%d' % tuple(map(ord, r[pos + 10:pos + 14]))
                )
            # skip over TYPE, CLASS, TTL, RDLENGTH, RDATA
            pos = pos + 8
            rdlength = (ord(r[pos]) << 8) + (ord(r[pos + 1]))
            pos = pos + 2 + rdlength
        return 0, None
    else:
        return 0, None


def unpack_ptr_reply(r):
    ancount = (ord(r[6]) << 8) + (ord(r[7]))
    # skip question, first name starts at 12,
    # this is followed by QTYPE and QCLASS
    pos = skip_name(r, 12) + 4
    if ancount:
        # we are looking very specifically for
        # an answer with TYPE=PTR, CLASS=IN (\000\014\000\001)
        for an in range(ancount):
            pos = skip_name(r, pos)
            if r[pos:pos + 4] == '\000\014\000\001':
                return (
                    unpack_ttl(r, pos + 4),
                    unpack_name(r, pos + 10)
                )
            # skip over TYPE, CLASS, TTL, RDLENGTH, RDATA
            pos = pos + 8
            rdlength = (ord(r[pos]) << 8) + (ord(r[pos + 1]))
            pos = pos + 2 + rdlength
        return 0, None
    else:
        return 0, None


# This is a UDP (datagram) resolver.

#
# It may be useful to implement a TCP resolver.  This would presumably
# give us more reliable behavior when things get too busy.  A TCP
# client would have to manage the connection carefully, since the
# server is allowed to close it at will (the RFC recommends closing
# after 2 minutes of idle time).
#
# Note also that the TCP client will have to prepend each request
# with a 2-byte length indicator (see rfc1035).
#

class resolver(asyncore.dispatcher):

    id = counter()

    def __init__(self, server='127.0.0.1'):
        asyncore.dispatcher.__init__(self)
        self.create_socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.server = server
        self.request_map = {}
        self.last_reap_time = int(time.time())      # reap every few minutes

    def writable(self):
        return 0

    def log(self, *args):
        pass

    def handle_close(self):
        self.log_info('closing!')
        self.close()

    def handle_error(self):
        # don't close the connection on error
        (file, fun, line), t, v, tbinfo = asyncore.compact_traceback()
        self.log_info(
            'Problem with DNS lookup (%s:%s %s)' % (t, v, tbinfo),
            'error')

    def get_id(self):
        return(self.id.as_long() % (1 << 16))

    def reap(self):
        # find DNS requests that have timed out
        now = int(time.time())
        if now - self.last_reap_time > 180:
            # reap every 3 minutes
            self.last_reap_time = now  # update before we forget
            for k, (host, unpack, callback, when) in self.request_map.items():
                if now - when > 180:
                    # over 3 minutes old
                    del self.request_map[k]
                    try:
                        # same code as in handle_read
                        callback(host, 0, None)  # timeout val is (0,None)
                    except Exception:
                        (file, fun, line), t, v, tbinfo = \
                            asyncore.compact_traceback()
                        self.log_info('%s %s %s' % (t, v, tbinfo), 'error')

    def resolve(self, host, callback):
        self.reap()  # first, get rid of old guys
        self.socket.sendto(
            fast_address_request(host, self.get_id()),
            (self.server, 53)
        )
        self.request_map[self.get_id()] = (
            host, unpack_address_reply, callback, int(time.time()))
        self.id.increment()

    def resolve_ptr(self, host, callback):
        self.reap()  # first, get rid of old guys
        ip = string.split(host, '.')
        ip.reverse()
        ip = string.join(ip, '.') + '.in-addr.arpa'
        self.socket.sendto(
            fast_ptr_request(ip, self.get_id()),
            (self.server, 53)
        )
        self.request_map[self.get_id()] = (
            host, unpack_ptr_reply, callback, int(time.time()))
        self.id.increment()

    def handle_read(self):
        reply, whence = self.socket.recvfrom(512)
        # for security reasons we may want to double-check
        # that <whence> is the server we sent the request to.
        id = (ord(reply[0]) << 8) + ord(reply[1])
        if id in self.request_map:
            host, unpack, callback, when = self.request_map[id]
            del self.request_map[id]
            ttl, answer = unpack(reply)
            try:
                callback(host, ttl, answer)
            except Exception:
                (file, fun, line), t, v, tbinfo = asyncore.compact_traceback()
                self.log_info('%s %s %s' % (t, v, tbinfo), 'error')


class rbl(resolver):

    def resolve_maps(self, host, callback):
        ip = string.split(host, '.')
        ip.reverse()
        ip = string.join(ip, '.') + '.rbl.maps.vix.com'
        self.socket.sendto(
            fast_ptr_request(ip, self.get_id()),
            (self.server, 53)
        )
        self.request_map[self.get_id()] = host, self.check_reply, callback
        self.id.increment()

    def check_reply(self, r):
        # we only need to check RCODE.
        rcode = (ord(r[3]) & 0xf)
        self.log_info('MAPS RBL; RCODE =%02x\n %s' % (rcode, repr(r)))
        return 0, rcode  # (ttl, answer)


class hooked_callback:
    def __init__(self, hook, callback):
        self.hook, self.callback = hook, callback

    def __call__(self, *args):
        self.hook(*args)
        self.callback(*args)


class caching_resolver(resolver):
    "Cache DNS queries.  Will need to honor the TTL value in the replies"

    def __init__(self, *args):
        resolver.__init__(self, *args)
        self.cache = {}
        self.forward_requests = counter()
        self.reverse_requests = counter()
        self.cache_hits = counter()

    def resolve(self, host, callback):
        self.forward_requests.increment()
        if host in self.cache:
            when, ttl, answer = self.cache[host]
            # ignore TTL for now
            callback(host, ttl, answer)
            self.cache_hits.increment()
        else:
            resolver.resolve(
                self,
                host,
                hooked_callback(
                    self.callback_hook,
                    callback
                )
            )

    def resolve_ptr(self, host, callback):
        self.reverse_requests.increment()
        if host in self.cache:
            when, ttl, answer = self.cache[host]
            # ignore TTL for now
            callback(host, ttl, answer)
            self.cache_hits.increment()
        else:
            resolver.resolve_ptr(
                self,
                host,
                hooked_callback(
                    self.callback_hook,
                    callback
                )
            )

    def callback_hook(self, host, ttl, answer):
        self.cache[host] = time.time(), ttl, answer

    SERVER_IDENT = 'Caching DNS Resolver (V%s)' % VERSION

    def status(self):
        from . import producers
        return producers.simple_producer(
            '<h2>%s</h2>' % self.SERVER_IDENT +
            '<br>Server: %s' % self.server +
            '<br>Cache Entries: %d' % len(self.cache) +
            '<br>Outstanding Requests: %d' % len(self.request_map) +
            '<br>Forward Requests: %s' % self.forward_requests +
            '<br>Reverse Requests: %s' % self.reverse_requests +
            '<br>Cache Hits: %s' % self.cache_hits
        )


if __name__ == '__main__':
    if len(sys.argv) == 1:
        print('usage: %s [-r] [-s <server_IP>] host [host ...]' % sys.argv[0])
        sys.exit(0)
    elif ('-s' in sys.argv):
        i = sys.argv.index('-s')
        server = sys.argv[i + 1]
        del sys.argv[i:i + 2]
    else:
        server = '127.0.0.1'

    if ('-r' in sys.argv):
        reverse = 1
        i = sys.argv.index('-r')
        del sys.argv[i]
    else:
        reverse = 0

    if ('-m' in sys.argv):
        maps = 1
        sys.argv.remove('-m')
    else:
        maps = 0

    if maps:
        r = rbl(server)
    else:
        r = caching_resolver(server)

    count = len(sys.argv) - 1

    def print_it(host, ttl, answer):
        global count
        print('%s: %s' % (host, answer))
        count = count - 1
        if not count:
            r.close()

    for host in sys.argv[1:]:
        if reverse:
            r.resolve_ptr(host, print_it)
        elif maps:
            r.resolve_maps(host, print_it)
        else:
            r.resolve(host, print_it)

            # hooked asyncore.loop()
    while asyncore.socket_map:
        asyncore.poll(30.0)
        print('requests outstanding: %d' % len(r.request_map))
