privrdns/privrdns/__init__.py

#!/usr/bin/env python
#
# Copyright 2019 John-Mark Gurney.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
#

__author__ = 'John-Mark Gurney'
__copyright__ = 'Copyright 2019 John-Mark Gurney.  All rights reserved.'
__license__ = '2-clause BSD license'
__version__ = '0.1.0.dev'

import asyncio
import dns
import socket
import unittest

from dnslib import DNSRecord, DNSQuestion, RR, QTYPE, A, digparser
from ntunnel import async_test, parsesockstr

class DNSCache(object):
	def __init__(self):
		self._data = {}

		# preload some values:
		self._data.update({
			('.', 'DS'): [ ('.', 'IN', 'DS', 20326, 8, 2,
			    bytes.fromhex('E06D44B80B8F1D39A95C0B0D7C65D08458E880409BBC683457104237C7F8EC8D'))],
		})

	def get(self, name, rtype):
		return self._data[(name, rtype)]

class Resolver(object):
	'''Resolve DNS names to records.

	Either serve it from the cache, or do a query.  If the query
	can be DNSSEC validated, do it via the privacy querier if not,
	via the trusted one.'''

def _writepkt(wrr, bts):
	wrr.write(len(bts).to_bytes(2, 'big'))
	wrr.write(bts)

async def _readpkt(rdr):
	nbytes = int.from_bytes(await rdr.readexactly(2), 'big')
	return await rdr.readexactly(nbytes)

class ServerResolver(object):
	'''Ask a specific server a question, and return the result.'''

	def __init__(self, rdrwrr):
		'''Pass in the reader and writer pair that is connected
		to the server.'''

		self._reader, self._writer = rdrwrr

		self._qs = {}

		self._replies = asyncio.create_task(self.procreplies(self._reader))

	def __del__(self):
		if self._replies:
			self._replies.cancel()
			self._replies = None

		self._writer.write_eof()

	def cleanup(self):
		while self._qs:
			i = self._qs.popitem()[1]
			i.set_exception(RuntimeError(
			    'server connection closed, '
			    'no response received'))

	async def procreplies(self, rdr):
		while True:
			# Get a response
			try:
				pkt = await _readpkt(self._reader)
			except asyncio.streams.IncompleteReadError:
				self.cleanup()
				return

			resp = DNSRecord.parse(pkt)

			# fetch where we should send it
			fut = self._qs.pop(resp.header.id, None)
			if fut is None:
				continue

			# send the result
			fut.set_result(resp.rr)

	async def resolve(self, q):
		# make DNS packet
		pkt = DNSRecord(questions=[q])
		pktbytes = pkt.pack()

		# Where the procreplies function will send the answer
		res = asyncio.Future()

		# Record it
		self._qs[pkt.header.id] = res

		# Send query
		_writepkt(self._writer, pktbytes)

		return await res

class DNSProc(asyncio.DatagramProtocol):
	def connection_made(self, transport):
		#print('cm')
		self.transport = transport

	def datagram_received(self, data, addr):
		#print('dr')
		pkt = DNSRecord.parse(data)
		#print(repr((pkt, addr)))

		d = pkt.reply()
		d.add_answer(RR("xxx.abc.com",QTYPE.A,rdata=A("1.2.3.4")))
		data = d.pack()
		self.transport.sendto(data, addr)

async def dnsprocessor(sockstr):
	proto, args = parsesockstr(sockstr)

	if proto == 'udp':
		loop = asyncio.get_event_loop()
		#print('pre-cde')
		trans, protocol = await loop.create_datagram_endpoint(DNSProc,
		    local_addr=(args.get('host', '127.0.0.1'), args['port']))
		#print('post-cde', repr((trans, protocol)), trans is protocol.transport)
	else:
		raise ValueError('unknown protocol: %s' % repr(proto))

def _asyncsockpair():
	'''Create a pair of sockets that are bound to each other.
	The function will return a tuple of two coroutine's, that
	each, when await'ed upon, will return the reader/writer pair.'''

	socka, sockb = socket.socketpair()

	return (asyncio.open_connection(sock=socka),
	    asyncio.open_connection(sock=sockb))

class Tests(unittest.TestCase):
	@async_test
	async def test_processdnsfailures(self):
		port = 5
		with self.assertRaises(ValueError):
			dnsproc = await dnsprocessor(
			    'tcp:host=127.0.0.1,port=%d' % port)
			#print('post-dns')

	@async_test
	async def test_processdns(self):
		# start the dns processor
		port = 38394
		dnsproc = await dnsprocessor(
			    'udp:host=127.0.0.1,port=%d' % port)

		# submit a query to it
		digproc = await asyncio.create_subprocess_exec('dig',
		    '@127.0.0.1', '-p', str(port), '+dnssec', '+time=1',
		    'www.funkthat.com.',
		    stdout=asyncio.subprocess.PIPE)

		# make sure it exits successfully
		self.assertEqual(await digproc.wait(), 0)

		stdout, stderr = await digproc.communicate()

		rep = list(digparser.DigParser(stdout))
		self.assertEqual(len(rep), 1)
		#print('x', repr(list(rep)))

	def test_cache(self):
		# test that the cache
		cache = DNSCache()

		# has the root trust anchors in it
		dsrs = cache.get('.', 'DS')

		self.assertEqual(dsrs, [ ('.', 'IN', 'DS', 20326, 8, 2,
		    bytes.fromhex('E06D44B80B8F1D39A95C0B0D7C65D08458E880409BBC683457104237C7F8EC8D')) ])

	@async_test
	async def xtest_query(self): # pragma: no cover
		# test that the query function
		q = await query_record('com.', 'DS')

	@async_test
	async def xtest_realresolver(self): # pragma: no cover
		# This is to test against a real resolver.
		host = 'gold'
		res = ServerResolver(await asyncio.open_connection(host, 53))

		resquestion = DNSQuestion('example.com')

	@async_test
	async def test_resolvernoanswer(self):
		client, server = _asyncsockpair()
		res = ServerResolver(await client)

		rdr, wrr = await server

		# constants
		resquestion = DNSQuestion('example.com')

		# start the query
		ans = asyncio.create_task(res.resolve(resquestion))

		# Fetch the question
		question = await _readpkt(rdr)
		dnsrec = DNSRecord.parse(question)

		# Make sure we got the correct question
		self.assertEqual(dnsrec.get_q(), resquestion)

		# close the connection
		wrr.write_eof()

		# make sure it raises an error
		with self.assertRaises(RuntimeError):
			await ans

	@async_test
	async def test_questionresolver(self):
		client, server = _asyncsockpair()
		res = ServerResolver(await client)

		rdr, wrr = await server

		# constants
		resquestion = DNSQuestion('example.com')
		resanswer = RR.fromZone('example.com. A 192.0.2.10\nexample.com. A 192.0.2.11')

		# start the query
		ans = asyncio.create_task(res.resolve(resquestion))

		# Fetch the question
		question = await _readpkt(rdr)
		dnsrec = DNSRecord.parse(question)

		# Make sure we got the correct question
		self.assertEqual(dnsrec.get_q(), resquestion)

		# Generate the reply
		rep = dnsrec.reply()
		rep.add_answer(*resanswer)
		repbytes = rep.pack()

		# Send the reply
		_writepkt(wrr, repbytes)
		wrr.write_eof()

		# veryify the received answer
		ans = await ans
		self.assertEqual(ans, resanswer)

	@async_test
	async def test_outoforderanswer(self):
		client, server = _asyncsockpair()
		res = ServerResolver(await client)

		rdr, wrr = await server

		# constants
		resquestiona = DNSQuestion('example.com')
		resquestionb = DNSQuestion('example.net')
		resanswera = RR.fromZone('example.com. A 192.0.2.10\nexample.com. A 192.0.2.11')
		resanswerb = RR.fromZone('example.net. A 192.0.2.20\nexample.net. A 192.0.2.21')

		# start the first query
		ansa = asyncio.create_task(res.resolve(resquestiona))

		# Fetch the first question
		question = await _readpkt(rdr)
		dnsreca = DNSRecord.parse(question)

		# Make sure we got the correct question
		self.assertEqual(dnsreca.get_q(), resquestiona)

		# start the second query (before first has been answered)
		ansb = asyncio.create_task(res.resolve(resquestionb))

		# Fetch the second question
		question = await _readpkt(rdr)
		dnsrecb = DNSRecord.parse(question)

		# Make sure we got the correct question
		self.assertEqual(dnsrecb.get_q(), resquestionb)

		# Generate the second reply
		rep = dnsrecb.reply()
		rep.add_answer(*resanswerb)
		repbytes = rep.pack()

		# Send the reply
		_writepkt(wrr, repbytes)

		# Generate the first reply
		rep = dnsreca.reply()
		rep.add_answer(*resanswera)
		repbytes = rep.pack()

		# Send the reply
		_writepkt(wrr, repbytes)

		# Send a second reply, and make sure it is ignored
		_writepkt(wrr, repbytes)

		# close the connection
		wrr.write_eof()

		# veryify the first received answer
		ansa = await ansa
		self.assertEqual(ansa, resanswera)

		# veryify the second received answer
		ansb = await ansb
		self.assertEqual(ansb, resanswerb)