import argparse import itertools import logging import os import random import signal import selectors import socket import sys import threading import time import dns.message import dns.rdatatype from pydnstest import scenario, mock_client from networking import InterfaceManager class TestServer: """ This simulates UDP DNS server returning scripted or mirror DNS responses. """ RETRIES_ON_BIND = 3 def __init__(self, test_scenario, addr_family, deckard_address=None, if_manager=None): """ Initialize server instance. """ self.thread = None self.srv_socks = [] self.client_socks = [] self.connections = [] self.active = False self.active_lock = threading.Lock() self.condition = threading.Condition() self.scenario = test_scenario self.scenario.deckard_address = deckard_address self.addr_map = [] self.start_iface = 2 self.cur_iface = self.start_iface self.addr_family = addr_family self.undefined_answers = 0 self.if_manager = if_manager def __del__(self): """ Cleanup after deletion. """ with self.active_lock: active = self.active if active: self.stop() def start(self): """ Synchronous start """ with self.active_lock: if self.active: raise Exception('TestServer already started') with self.active_lock: self.active = True self._bind_sockets() def stop(self): """ Stop socket server operation. """ with self.active_lock: self.active = False if self.thread: self.thread.join() for conn in self.connections: conn.close() for srv_sock in self.srv_socks: srv_sock.close() for client_sock in self.client_socks: client_sock.close() self.client_socks = [] self.srv_socks = [] self.connections = [] self.scenario = None def address(self): """ Returns opened sockets list """ addrlist = [] for s in self.srv_socks: addrlist.append(s.getsockname()) return addrlist def handle_query(self, client): """ Receive query from client socket and send an answer. Returns: True if client socket should be closed by caller False if client socket should be kept open """ log = logging.getLogger('pydnstest.testserver.handle_query') server_addr = client.getsockname()[0] query, client_addr = mock_client.recvfrom_msg(client) if query is None: return False log.debug('server %s received query from %s: %s', server_addr, client_addr, query) message = self.scenario.reply(query, server_addr) if not message: log.debug('ignoring') return True elif isinstance(message, scenario.DNSReplyServfail): self.undefined_answers += 1 self.scenario.current_step.log.error( 'server %s has no response for question %s, answering with SERVFAIL', server_addr, '; '.join([str(rr) for rr in query.question])) else: log.debug('response: %s', message) mock_client.sendto_msg(client, message.to_wire(), client_addr) return True def query_io(self): """ Main server process """ self.undefined_answers = 0 with self.active_lock: if not self.active: raise Exception("[query_io] Test server not active") while True: with self.condition: self.condition.notify() with self.active_lock: if not self.active: break objects = self.srv_socks + self.connections sel = selectors.DefaultSelector() for obj in objects: sel.register(obj, selectors.EVENT_READ) items = sel.select(0.1) for key, event in items: sock = key.fileobj if event & selectors.EVENT_READ: if sock in self.srv_socks: if sock.proto == socket.IPPROTO_TCP: conn, _ = sock.accept() self.connections.append(conn) else: self.handle_query(sock) elif sock in self.connections: if not self.handle_query(sock): sock.close() self.connections.remove(sock) else: raise Exception( "[query_io] Socket IO internal error {}, exit" .format(sock.getsockname())) else: raise Exception("[query_io] Socket IO error {}, exit" .format(sock.getsockname())) def start_srv(self, address, family, proto=socket.IPPROTO_UDP): """ Starts listening thread if necessary """ assert address assert address[0] # host assert address[1] # port assert family assert proto if family == socket.AF_INET6: if not socket.has_ipv6: raise NotImplementedError("[start_srv] IPv6 is not supported by socket {0}" .format(socket)) elif family != socket.AF_INET: raise NotImplementedError("[start_srv] unsupported protocol family {0}".format(family)) if proto == socket.IPPROTO_TCP: socktype = socket.SOCK_STREAM elif proto == socket.IPPROTO_UDP: socktype = socket.SOCK_DGRAM else: raise NotImplementedError("[start_srv] unsupported protocol {0}".format(proto)) if self.thread is None: self.thread = threading.Thread(target=self.query_io) self.thread.start() with self.condition: self.condition.wait() for srv_sock in self.srv_socks: if (srv_sock.family == family and srv_sock.getsockname()[:2] == address and srv_sock.proto == proto): return sock = socket.socket(family, socktype, proto) sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) # Add address to interface when running from Deckard if self.if_manager is not None: if address[0] not in self.if_manager.added_addresses: self.if_manager.add_address(address[0]) # A lot of addresses are added to the interface while runnning from Deckard in # the small amount of time which caused ocassional hiccups while binding to them # right afterwards in testing. Therefore, we retry a few times. final_ex = None for i in range(self.RETRIES_ON_BIND): try: sock.bind(address) break except OSError as ex: # Exponential backoff time.sleep((2 ** i) + random.random()) final_ex = ex continue else: print(final_ex, address) raise final_ex if proto == socket.IPPROTO_TCP: sock.listen(5) self.srv_socks.append(sock) def _bind_sockets(self): """ Bind test server to port 53 on all addresses referenced by test scenario. """ # Bind to test servers for r in self.scenario.ranges: for addr in r.addresses: family = socket.AF_INET6 if ':' in addr else socket.AF_INET self.start_srv((addr, 53), family) self.start_srv((addr, 53), family, proto=socket.IPPROTO_TCP) # Bind addresses in ad-hoc REPLYs for s in self.scenario.steps: if s.type == 'REPLY': reply = s.data[0].message for rr in itertools.chain(reply.answer, reply.additional, reply.question, reply.authority): for rd in rr: if rd.rdtype == dns.rdatatype.A: self.start_srv((rd.address, 53), socket.AF_INET) self.start_srv((rd.address, 53), socket.AF_INET, proto=socket.IPPROTO_TCP) elif rd.rdtype == dns.rdatatype.AAAA: self.start_srv((rd.address, 53), socket.AF_INET6) self.start_srv((rd.address, 53), socket.AF_INET6, proto=socket.IPPROTO_TCP) def play(self, subject_addr): self.scenario.play({'': (subject_addr, 53)}) def empty_test_case(): """ Return (scenario, config) pair which answers to any query on 127.0.0.10. """ # Mirror server empty_test_path = os.path.dirname(os.path.realpath(__file__)) + "/empty.rpl" test_config = {'ROOT_ADDR': '127.0.0.10', '_SOCKET_FAMILY': socket.AF_INET} return scenario.parse_file(empty_test_path)[0], test_config def standalone_self_test(): """ Self-test code Usage: unshare -rn $PYTHON -m pydnstest.testserver --help """ logging.basicConfig(level=logging.DEBUG) argparser = argparse.ArgumentParser() argparser.add_argument('--scenario', help='absolute path to test scenario', required=False) argparser.add_argument('--step', help='step # in the scenario (default: first)', required=False, type=int) args = argparser.parse_args() if args.scenario: test_scenario, test_config_text = scenario.parse_file(args.scenario) test_config = scenario.parse_config(test_config_text, True, os.getcwd()) else: test_scenario, test_config = empty_test_case() if args.step: for step in test_scenario.steps: if step.id == args.step: test_scenario.current_step = step if not test_scenario.current_step: raise ValueError('step ID %s not found in scenario' % args.step) else: test_scenario.current_step = test_scenario.steps[0] if_manager = InterfaceManager(interface="testserver") server = TestServer(test_scenario, test_config['_SOCKET_FAMILY'], if_manager=if_manager) server.start() logging.info("[==========] Mirror server running at %s", server.address()) def kill(signum, frame): # pylint: disable=unused-argument logging.info("[==========] Shutdown.") server.stop() sys.exit(128 + signum) signal.signal(signal.SIGINT, kill) signal.signal(signal.SIGTERM, kill) while True: time.sleep(0.5) if __name__ == '__main__': # this is done to avoid creating global variables standalone_self_test()