def test_parse_ns_question(self): expected = DnsPacket(id_=2, is_authoritative=False, is_response=False, opcode=0, is_truncated=False, is_recursion_desired=True, is_recursion_available=False, question=(b'e1.ru', QueryType.NS), answers=(), has_error=False) self._test_packet(expected, 1)
def test_parse_ptr_question(self): expected = DnsPacket(id_=1, is_authoritative=False, is_response=False, opcode=0, is_truncated=False, is_recursion_desired=True, is_recursion_available=False, question=(b'1.1.168.192.in-addr.arpa', QueryType.PTR), answers=(), has_error=False) self._test_packet(expected, 0)
def _test_packet(self, expected, index): packet = DnsPacket.from_bytes(DnsPacketTest.packets[index]) self.assertEqual(packet.id, expected.id) self.assertEqual(packet.has_error, expected.has_error) self.assertEqual(packet.is_authoritative, expected.is_authoritative) self.assertEqual(packet.is_recursion_desired, expected.is_recursion_desired) self.assertEqual(packet.is_recursion_available, expected.is_recursion_available) self.assertEqual(packet.is_response, expected.is_response) self.assertEqual(packet.is_truncated, expected.is_truncated) self.assertEqual(packet.opcode, expected.opcode) self.assertEqual(packet.question, expected.question) self.assertEqual(packet.answers, expected.answers)
def test_parse_ns_answers(self): answers = [(b'e1.ru', QueryType.NS, 300, b'\x03ns2\x03ngs\xc0\x0f'), (b'e1.ru', QueryType.NS, 300, b"\x02ns\xc0'"), (b'e1.ru', QueryType.NS, 300, b'\x03ns1\xc0\x0c'), (b'e1.ru', QueryType.NS, 300, b'\x03ns2\xc0\x0c')] expected = DnsPacket(id_=2, is_authoritative=False, is_response=True, opcode=0, is_truncated=False, is_recursion_available=True, is_recursion_desired=True, question=(b'e1.ru', QueryType.NS), answers=tuple(answers), has_error=False) self._test_packet(expected, 2)
def on_connection(self, sock, data, addr): try: print('>> Data received') packet = DnsPacket.from_bytes(data) name, query_type = packet.question print('>> Question: ', query_type, name) cache_key = int(query_type), name if cache_key in self.cache: print('>> Found fresh record in cache') sock.sendto(self.cache[cache_key], addr) else: print('>> No fresh record in cache') origin_packet = self.forward_question(packet) if not origin_packet.answers: print('>> Origin has no answers') return min_ttl = min(origin_packet.answers, key=lambda x: x[2])[2] self.cache.add(cache_key, bytes(origin_packet), min_ttl) sock.sendto(bytes(origin_packet), addr) print('>> Answered') except: print('>> Error with packet: {}\n{}'.format( repr(data), format_exc()))
def forward_question(self, packet): sock = socket(AF_INET, SOCK_DGRAM) sock.settimeout(12) sock.sendto(bytes(packet), self.origin) data, addr = sock.recvfrom(4096) return DnsPacket.from_bytes(data)
def test_serialize(self): for packet in DnsPacketTest.packets: parsed = DnsPacket.from_bytes(packet) serialized = bytes(parsed) self.assertEqual(serialized, packet)
def test_serialize_ptr_question(self): packet = DnsPacketTest.packets[0] parsed = DnsPacket.from_bytes(packet) serialized = bytes(parsed) self.assertEqual(len(serialized), len(packet)) self.assertEqual(serialized, packet)