Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #!/usr/bin/env python3
- """
- Forward DNS query (in UDP) to a set of servers in parallel, then send back
- the fastest reply and ignore others. TCP queries are directly forward to a
- fixed server.
- """
- import asyncio
- import socket
- LISTEN = ('127.0.0.1', 5353)
- UDP_SERVERS = [('127.0.0.1', port) for port in range(8130, 8135)]
- TCP_SERVER = ('8.8.8.8', 53)
- UDP_MAX_WAIT = 5
- TCP_TIMEOUT = 10
- class UDPServerProtocol:
- def connection_made(self, transport):
- self.transport = transport
- self.loop = asyncio.get_event_loop()
- def datagram_received(self, data, addr):
- print(f'receive {len(data)} bytes from {addr}')
- sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
- sock.setblocking(False)
- for remote in UDP_SERVERS:
- sock.sendto(data, remote)
- coro = self.loop.sock_recv(sock, 8192)
- task = self.loop.create_task(asyncio.wait_for(coro, UDP_MAX_WAIT))
- def done(task):
- if task.exception() is not None:
- if isinstance(task.exception(), asyncio.TimeoutError):
- print('no any reply from all servers')
- else:
- print(f'error on receiving udp: {task.exception()}')
- else:
- resp = task.result()
- print(f'receive response {len(resp)} bytes')
- self.transport.sendto(resp, addr)
- sock.close()
- task.add_done_callback(done)
- async def tcp_forwarder(local_r, local_w):
- peername = local_w.get_extra_info('peername')
- print(f'tcp connected from {peername}')
- remote_r, remote_w = await asyncio.open_connection(*TCP_SERVER)
- #print(f'connected to remote')
- pipings = [piping(local_r, remote_w), piping(remote_r, local_w)]
- await asyncio.wait(pipings, return_when=asyncio.FIRST_EXCEPTION)
- local_w.close()
- remote_w.close()
- #print(f'tcp close')
- async def piping(reader, writer):
- while True:
- data = await asyncio.wait_for(reader.read(8196), TCP_TIMEOUT)
- if not data:
- writer.write_eof()
- return
- writer.write(data)
- await writer.drain()
- def main():
- loop = asyncio.get_event_loop()
- listen = loop.create_datagram_endpoint(UDPServerProtocol, LISTEN)
- udp, proto = loop.run_until_complete(listen)
- listen = asyncio.start_server(tcp_forwarder, *LISTEN)
- tcp = loop.run_until_complete(listen)
- print('running')
- try:
- loop.run_forever()
- except KeyboardInterrupt:
- pass
- udp.close()
- tcp.close()
- loop.close()
- if __name__ == '__main__':
- main()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement