# -*- Mode: Python; tab-width: 4 -*-
# 	$Id: asyndns.py,v 1.2 1996/02/07 01:28:52 rushing Exp $
#	Author: Sam Rushing <rushing@nightmare.com>

import dnslib
import npdns
import socket
import string

import asyncore
import asynchat

error = 'async dns library error'

# need an optional limit on the number of outstanding requests.

class dns_client (asyncore.dispatcher):
	def __init__ (self, server=None):
		self.server = server
		# maps request ID's with done_funs.
		self.reply_map = {}
		asyncore.dispatcher.__init__ (self)
		self.create_socket (socket.AF_INET, socket.SOCK_DGRAM)

	def send_request (self, q, server):
		if not server:
			if not self.server:
				raise ValueError, 'no default server'
			else:
				server = self.server
		#print 'sending %d bytes to %s' % (len(q), server)
		self.sendto (q, (server, 53))

	def get_host_by_name (self, name, done_fun, server=None):
		q, id = npdns.build_address_request (name)
		self.reply_map[id] = done_fun
		self.send_request (q, server)
		return id

	def get_host_by_addr (self, addr, done_fun, server=None):
		q, id = npdns.build_ptr_request (addr)
		self.reply_map[id] = done_fun
		self.send_request (q, server)
		return id

	def handle_read (self):
		# Maximum DNS message size is 512.  does the UDP layer guarantee
		# we'll always get the whole message?
		reply, addr = self.recvfrom (512)
		info, ilen = npdns.Message.unpack (reply)
		h = info['header']
		id = h['id']
		#dnslib.dumpM (dnslib.Munpacker (reply))
		if self.reply_map.has_key (id):
			fun = self.reply_map[id]
		else:
			print 'unknown request id(%d) from %s' % (id, repr(addr))
			return
		#
		if h['rcode'] == 0:
			if h['ancount'] > 0:
				fun (
					id,
					info['questions'][0]['qname'],
					info['resource records'][0][0]['rdata']
					)
				del self.reply_map[id]
			else:
				fun (
					id,
					info['questions'][0]['qname'],
					'*** no answer ***'
					)
		else:
			fun (
				id,
				info['questions'][0]['qname'],
				'*** %s ***' % rcode_descriptions[h['rcode']]
				)
		
rcode_descriptions = [
	'no error',
	'format error',
	'server failure',
	'name error',
	'not implemented',
	'refused'
	]

# This is a 'different' type of name server - it sits on port 8053,
# reading either a FQDN or dotted-quad IP address, one to a line, and
# then sends the answers back asynchronously.

class demo_server (asyncore.dispatcher):
	def __init__ (self, resolver, port=8053):
		self.resolver = resolver
		asyncore.dispatcher.__init__ (self)
		self.create_socket (socket.AF_INET, socket.SOCK_STREAM)
		self.bind ('', port)

	def go (self):
		self.listen (10)

	def handle_accept (self):
		conn, addr = self.accept()
		demo (conn, self.resolver)

import regex
ip_reg = regex.compile ('[0-9]+\\.[0-9]+\\.[0-9]+\\.[0-9]+')

class demo (asynchat.async_chat):
	def __init__ (self, socket, resolver):
		self.resolver = resolver
		asynchat.async_chat.__init__ (self, socket)
		self.set_terminator ('\r\n')
		self.data = ''

	def collect_incoming_data (self, data):
		self.data = self.data + data

	def resolver_done_fun (self, id, question, answer):
		self.send ('%s = %s\n' %  (question, answer))

	def found_terminator (self):
		line = self.data
		self.data = ''
		if ip_reg.match (line) == len(line):
			self.resolver.get_host_by_addr (line, self.resolver_done_fun)
		else:
			self.resolver.get_host_by_name (line, self.resolver_done_fun)

def test (server='205.160.176.5', user_port=8053):
	r = dns_client(server)
	d = demo_server (r, user_port)
	d.go()

import sys

if __name__ == '__main__':
	import asyncore
	if len(sys.argv) > 1:
		if len (sys.argv) > 2:
			test (sys.argv[1], string.atoi (sys.argv[2]))
		else:
			test (sys.argv[1])
	else:
		test()
	asyncore.loop()
