testserver.py 10.5 KB
Newer Older
Ivana Krumlova's avatar
Ivana Krumlova committed
1 2 3
from __future__ import absolute_import
from __future__ import print_function

Marek Vavruša's avatar
Marek Vavruša committed
4 5 6 7 8 9 10 11
import threading
import select
import socket
import os
import time
import dns.message
import dns.rdatatype
import itertools
12 13
import struct
import binascii
Ivana Krumlova's avatar
Ivana Krumlova committed
14
from pydnstest.dprint import dprint
Marek Vavruša's avatar
Marek Vavruša committed
15

Petr Špaček's avatar
Petr Špaček committed
16 17

def recvfrom_msg(stream, raw=False):
18 19 20 21 22 23 24
    """
    Receive DNS message from TCP/UDP socket.

    Returns:
        if raw == False: (DNS message object, peer address)
        if raw == True: (blob, peer address)
    """
Ivana Krumlova's avatar
Ivana Krumlova committed
25
    if stream.type & socket.SOCK_DGRAM:
26
        data, addr = stream.recvfrom(4096)
Ivana Krumlova's avatar
Ivana Krumlova committed
27
    elif stream.type & socket.SOCK_STREAM:
28 29 30
        data = stream.recv(2)
        if len(data) == 0:
            return None, None
Petr Špaček's avatar
Petr Špaček committed
31
        msg_len = struct.unpack_from("!H", data)[0]
Ivana Krumlova's avatar
Ivana Krumlova committed
32
        data = b""
33 34 35 36 37 38
        received = 0
        while received < msg_len:
            next_chunk = stream.recv(4096)
            if len(next_chunk) == 0:
                return None, None
            data += next_chunk
Petr Špaček's avatar
Petr Špaček committed
39
            received += len(next_chunk)
40 41
        addr = stream.getpeername()[0]
    else:
Ivana Krumlova's avatar
Ivana Krumlova committed
42
        raise NotImplementedError("[recvfrom_msg]: unknown socket type '%i'" % stream.type)
43
    if not raw:
44
        data = dns.message.from_wire(data, one_rr_per_rrset=True)
45
    return data, addr
Marek Vavruša's avatar
Marek Vavruša committed
46

Petr Špaček's avatar
Petr Špaček committed
47 48

def sendto_msg(stream, message, addr=None):
49
    """ Send DNS/UDP/TCP message. """
Marek Vavruša's avatar
Marek Vavruša committed
50
    try:
Ivana Krumlova's avatar
Ivana Krumlova committed
51
        if stream.type & socket.SOCK_DGRAM:
52 53 54 55
            if addr is None:
                stream.send(message)
            else:
                stream.sendto(message, addr)
Ivana Krumlova's avatar
Ivana Krumlova committed
56
        elif stream.type & socket.SOCK_STREAM:
Petr Špaček's avatar
Petr Špaček committed
57
            data = struct.pack("!H", len(message)) + message
58 59
            stream.send(data)
        else:
Ivana Krumlova's avatar
Ivana Krumlova committed
60
            assert False, "[sendto_msg]: unknown socket type '%i'" % stream.type
Petr Špaček's avatar
Petr Špaček committed
61
    except:  # Failure to respond is OK, resolver should recover
Marek Vavruša's avatar
Marek Vavruša committed
62 63
        pass

Petr Špaček's avatar
Petr Špaček committed
64

Marek Vavruša's avatar
Marek Vavruša committed
65 66 67 68 69 70 71
def get_local_addr_str(family, iface):
    """ Returns pattern string for localhost address  """
    if family == socket.AF_INET:
        addr_local_pattern = "127.0.0.{}"
    elif family == socket.AF_INET6:
        addr_local_pattern = "fd00::5357:5f{:02X}"
    else:
Ivana Krumlova's avatar
Ivana Krumlova committed
72
        raise NotImplementedError("[get_local_addr_str] family not supported '%i'" % family)
Marek Vavruša's avatar
Marek Vavruša committed
73 74
    return addr_local_pattern.format(iface)

Petr Špaček's avatar
Petr Špaček committed
75

Marek Vavruša's avatar
Marek Vavruša committed
76 77
class AddrMapInfo:
    """ Saves mapping info between adresses from rpl and cwrap adresses """
Petr Špaček's avatar
Petr Špaček committed
78

Marek Vavruša's avatar
Marek Vavruša committed
79
    def __init__(self, family, local, external):
Petr Špaček's avatar
Petr Špaček committed
80 81 82 83
        self.family = family
        self.local = local
        self.external = external

Marek Vavruša's avatar
Marek Vavruša committed
84 85 86 87

class TestServer:
    """ This simulates UDP DNS server returning scripted or mirror DNS responses. """

88
    def __init__(self, scenario, config, d_iface):
Marek Vavruša's avatar
Marek Vavruša committed
89 90 91 92
        """ Initialize server instance. """
        self.thread = None
        self.srv_socks = []
        self.client_socks = []
93
        self.connections = []
Marek Vavruša's avatar
Marek Vavruša committed
94 95 96 97 98 99 100
        self.active = False
        self.scenario = scenario
        self.config = config
        self.addr_map = []
        self.start_iface = 2
        self.cur_iface = self.start_iface
        self.kroot_local = None
101
        self.addr_family = None
Marek Vavruša's avatar
Marek Vavruša committed
102
        self.default_iface = d_iface
103
        self.set_initial_address()
Marek Vavruša's avatar
Marek Vavruša committed
104 105 106 107 108 109

    def __del__(self):
        """ Cleanup after deletion. """
        if self.active is True:
            self.stop()

Petr Špaček's avatar
Petr Špaček committed
110
    def start(self, port=53):
Marek Vavruša's avatar
Marek Vavruša committed
111 112 113 114
        """ Synchronous start """
        if self.active is True:
            raise Exception('TestServer already started')
        self.active = True
115 116
        self.addr, _ = self.start_srv((self.kroot_local, port), self.addr_family)
        self.start_srv(self.addr, self.addr_family, socket.IPPROTO_TCP)
Marek Vavruša's avatar
Marek Vavruša committed
117 118 119 120

    def stop(self):
        """ Stop socket server operation. """
        self.active = False
121 122
        if self.thread:
            self.thread.join()
123 124
        for conn in self.connections:
            conn.close()
Marek Vavruša's avatar
Marek Vavruša committed
125 126 127 128 129 130
        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 = []
131
        self.connections = []
Marek Vavruša's avatar
Marek Vavruša committed
132 133
        self.scenario = None

Petr Špaček's avatar
Petr Špaček committed
134
    def check_family(self, addr, family):
135 136
        """ Determines if address matches family """
        test_addr = None
Marek Vavruša's avatar
Marek Vavruša committed
137 138
        try:
            n = socket.inet_pton(family, addr)
139
            test_addr = socket.inet_ntop(family, n)
Marek Vavruša's avatar
Marek Vavruša committed
140
        except socket.error:
141 142
            return False
        return True
Marek Vavruša's avatar
Marek Vavruša committed
143

144 145
    def set_initial_address(self):
        """ Set address for starting thread """
Marek Vavruša's avatar
Marek Vavruša committed
146
        if self.config is None:
147 148
            self.addr_family = socket.AF_INET
            self.kroot_local = get_local_addr_str(self.addr_family, self.default_iface)
Marek Vavruša's avatar
Marek Vavruša committed
149
            return
150
        # Default address is localhost
Marek Vavruša's avatar
Marek Vavruša committed
151 152 153 154 155
        kroot_addr = None
        for k, v in self.config:
            if k == 'stub-addr':
                kroot_addr = v
        if kroot_addr is not None:
Petr Špaček's avatar
Petr Špaček committed
156
            if self.check_family(kroot_addr, socket.AF_INET):
157
                self.addr_family = socket.AF_INET
158
                self.kroot_local = kroot_addr
Petr Špaček's avatar
Petr Špaček committed
159
            elif self.check_family(kroot_addr, socket.AF_INET6):
160
                self.addr_family = socket.AF_INET6
161 162
                self.kroot_local = kroot_addr
        else:
163 164
            self.addr_family = socket.AF_INET
            self.kroot_local = get_local_addr_str(self.addr_family, self.default_iface)
165

Marek Vavruša's avatar
Marek Vavruša committed
166 167
    def address(self):
        """ Returns opened sockets list """
Petr Špaček's avatar
Petr Špaček committed
168
        addrlist = []
Marek Vavruša's avatar
Marek Vavruša committed
169
        for s in self.srv_socks:
Petr Špaček's avatar
Petr Špaček committed
170 171
            addrlist.append(s.getsockname())
        return addrlist
Marek Vavruša's avatar
Marek Vavruša committed
172 173

    def handle_query(self, client):
174 175 176 177 178 179 180
        """
        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
        """
181
        client_address = client.getsockname()[0]
Marek Vavruša's avatar
Marek Vavruša committed
182 183 184
        query, addr = recvfrom_msg(client)
        if query is None:
            return False
Petr Špaček's avatar
Petr Špaček committed
185
        dprint("[ handle_query ]", "%s incoming query from %s\n%s" % (client_address, addr, query))
Marek Vavruša's avatar
Marek Vavruša committed
186 187 188 189 190 191
        response = dns.message.make_response(query)
        is_raw_data = False
        if self.scenario is not None:
            response, is_raw_data = self.scenario.reply(query, client_address)
        if response:
            if is_raw_data is False:
Petr Špaček's avatar
Petr Špaček committed
192 193
                data_to_wire = response.to_wire(max_size=65535)
                dprint("[ handle_query ]", "response\n%s" % response)
Marek Vavruša's avatar
Marek Vavruša committed
194
            else:
195
                data_to_wire = response
Petr Špaček's avatar
Petr Špaček committed
196
                dprint("[ handle_query ]", "raw response found")
Marek Vavruša's avatar
Marek Vavruša committed
197 198
        else:
            response = dns.message.make_response(query)
199
            response.set_rcode(dns.rcode.SERVFAIL)
200
            data_to_wire = response.to_wire()
Petr Špaček's avatar
Petr Špaček committed
201
            dprint("[ handle_query ]", "response failed, SERVFAIL")
202 203 204

        sendto_msg(client, data_to_wire, addr)
        return True
Marek Vavruša's avatar
Marek Vavruša committed
205 206 207 208 209 210

    def query_io(self):
        """ Main server process """
        if self.active is False:
            raise Exception("[query_io] Test server not active")
        while self.active is True:
Petr Špaček's avatar
Petr Špaček committed
211 212 213 214 215 216 217 218 219 220 221 222 223 224
            objects = self.srv_socks + self.connections
            to_read, _, to_error = select.select(objects, [], objects, 0.1)
            for sock in to_read:
                if sock in self.srv_socks:
                    if (sock.proto == socket.IPPROTO_TCP):
                        conn, addr = 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:
Petr Špaček's avatar
Petr Špaček committed
225 226
                    raise Exception(
                        "[query_io] Socket IO internal error {}, exit".format(sock.getsockname()))
Petr Špaček's avatar
Petr Špaček committed
227 228
            for sock in to_error:
                raise Exception("[query_io] Socket IO error {}, exit".format(sock.getsockname()))
Marek Vavruša's avatar
Marek Vavruša committed
229

Petr Špaček's avatar
Petr Špaček committed
230
    def start_srv(self, address=None, family=socket.AF_INET, proto=socket.IPPROTO_UDP):
Marek Vavruša's avatar
Marek Vavruša committed
231 232 233 234
        """ Starts listening thread if necessary """
        if family == None:
            family = socket.AF_INET
        if family == socket.AF_INET:
235 236
            if address[0] is None:
                address = (get_local_addr_str(family, self.default_iface), 53)
Marek Vavruša's avatar
Marek Vavruša committed
237 238
        elif family == socket.AF_INET6:
            if socket.has_ipv6 is not True:
Ivana Krumlova's avatar
Ivana Krumlova committed
239 240
                raise NotImplementedError("[start_srv] IPv6 is not supported by socket {0}"
                                          .format(socket))
241 242
            if address[0] is None:
                address = (get_local_addr_str(family, self.default_iface), 53)
243
        else:
Petr Špaček's avatar
Petr Špaček committed
244
            raise NotImplementedError("[start_srv] unsupported protocol family {0}".format(family))
245 246 247 248 249 250 251

        if proto == None:
            proto = socket.IPPROTO_UDP
        if proto == socket.IPPROTO_TCP:
            socktype = socket.SOCK_STREAM
        elif proto == socket.IPPROTO_UDP:
            socktype = socket.SOCK_DGRAM
Marek Vavruša's avatar
Marek Vavruša committed
252
        else:
Ivana Krumlova's avatar
Ivana Krumlova committed
253
            raise NotImplementedError("[start_srv] unsupported protocol {0}".format(proto))
254

Marek Vavruša's avatar
Marek Vavruša committed
255 256 257 258 259
        if (self.thread is None):
            self.thread = threading.Thread(target=self.query_io)
            self.thread.start()

        for srv_sock in self.srv_socks:
Petr Špaček's avatar
Petr Špaček committed
260 261 262
            if (srv_sock.family == family
                    and srv_sock.getsockname() == address
                    and srv_sock.proto == proto):
Marek Vavruša's avatar
Marek Vavruša committed
263
                return srv_sock.getsockname()
264 265

        sock = socket.socket(family, socktype, proto)
266
        sock.bind(address)
Marek Vavruša's avatar
Marek Vavruša committed
267
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
268 269
        if proto == socket.IPPROTO_TCP:
            sock.listen(5)
Marek Vavruša's avatar
Marek Vavruša committed
270 271
        self.srv_socks.append(sock)
        sockname = sock.getsockname()
272
        return sockname, proto
Marek Vavruša's avatar
Marek Vavruša committed
273

274
    def play(self, subject_addr):
275 276 277
        sockfamily = socket.AF_INET
        if self.scenario.force_ipv6 == True:
            sockfamily = socket.AF_INET6
278
        paddr = get_local_addr_str(sockfamily, subject_addr)
Marek Vavrusa's avatar
Marek Vavrusa committed
279
        self.scenario.play(sockfamily, {'': (paddr, 53)})
Marek Vavruša's avatar
Marek Vavruša committed
280 281 282

if __name__ == '__main__':
    # Self-test code
Ivana Krumlova's avatar
Ivana Krumlova committed
283
    # Usage: $PYTHON -m pydnstest.testserver
Marek Vavruša's avatar
Marek Vavruša committed
284 285 286
    DEFAULT_IFACE = 0
    CHILD_IFACE = 0
    if "SOCKET_WRAPPER_DEFAULT_IFACE" in os.environ:
Petr Špaček's avatar
Petr Špaček committed
287
        DEFAULT_IFACE = int(os.environ["SOCKET_WRAPPER_DEFAULT_IFACE"])
Petr Špaček's avatar
Petr Špaček committed
288
    if DEFAULT_IFACE < 2 or DEFAULT_IFACE > 254:
Marek Vavruša's avatar
Marek Vavruša committed
289
        DEFAULT_IFACE = 10
Petr Špaček's avatar
Petr Špaček committed
290
        os.environ["SOCKET_WRAPPER_DEFAULT_IFACE"] = "{}".format(DEFAULT_IFACE)
Marek Vavruša's avatar
Marek Vavruša committed
291
    # Mirror server
Petr Špaček's avatar
Petr Špaček committed
292
    server = TestServer(None, None, DEFAULT_IFACE)
Marek Vavruša's avatar
Marek Vavruša committed
293
    server.start()
Ivana Krumlova's avatar
Ivana Krumlova committed
294
    print("[==========] Mirror server running at", server.address())
Marek Vavruša's avatar
Marek Vavruša committed
295 296
    try:
        while True:
Ivana Krumlova's avatar
Ivana Krumlova committed
297
            time.sleep(0.5)
Marek Vavruša's avatar
Marek Vavruša committed
298
    except KeyboardInterrupt:
Ivana Krumlova's avatar
Ivana Krumlova committed
299
        print("[==========] Shutdown.")
Marek Vavruša's avatar
Marek Vavruša committed
300 301
        pass
    server.stop()