SHARE
TWEET

Untitled

a guest Jul 17th, 2017 40 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #!/usr/bin/env python3
  2. """
  3. Forward DNS query (in UDP) to a set of servers in parallel, then send back
  4. the fastest reply and ignore others. TCP queries are directly forward to a
  5. fixed server.
  6. """
  7. import asyncio
  8. import socket
  9.  
  10.  
  11. LISTEN = ('127.0.0.1', 5353)
  12. UDP_SERVERS = [('127.0.0.1', port) for port in range(8130, 8135)]
  13. TCP_SERVER = ('8.8.8.8', 53)
  14. UDP_MAX_WAIT = 5
  15. TCP_TIMEOUT = 10
  16.  
  17. class UDPServerProtocol:
  18.     def connection_made(self, transport):
  19.         self.transport = transport
  20.         self.loop = asyncio.get_event_loop()
  21.  
  22.     def datagram_received(self, data, addr):
  23.         print(f'receive {len(data)} bytes from {addr}')
  24.         sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  25.         sock.setblocking(False)
  26.         for remote in UDP_SERVERS:
  27.             sock.sendto(data, remote)
  28.         coro = self.loop.sock_recv(sock, 8192)
  29.         task = self.loop.create_task(asyncio.wait_for(coro, UDP_MAX_WAIT))
  30.         def done(task):
  31.             if task.exception() is not None:
  32.                 if isinstance(task.exception(), asyncio.TimeoutError):
  33.                     print('no any reply from all servers')
  34.                 else:
  35.                     print(f'error on receiving udp: {task.exception()}')
  36.             else:
  37.                 resp = task.result()
  38.                 print(f'receive response {len(resp)} bytes')
  39.                 self.transport.sendto(resp, addr)
  40.             sock.close()
  41.         task.add_done_callback(done)
  42.  
  43.  
  44. async def tcp_forwarder(local_r, local_w):
  45.     peername = local_w.get_extra_info('peername')
  46.     print(f'tcp connected from {peername}')
  47.     remote_r, remote_w = await asyncio.open_connection(*TCP_SERVER)
  48.     #print(f'connected to remote')
  49.     pipings = [piping(local_r, remote_w), piping(remote_r, local_w)]
  50.     await asyncio.wait(pipings, return_when=asyncio.FIRST_EXCEPTION)
  51.     local_w.close()
  52.     remote_w.close()
  53.     #print(f'tcp close')
  54.  
  55.  
  56. async def piping(reader, writer):
  57.     while True:
  58.         data = await asyncio.wait_for(reader.read(8196), TCP_TIMEOUT)
  59.         if not data:
  60.             writer.write_eof()
  61.             return
  62.         writer.write(data)
  63.         await writer.drain()
  64.  
  65.  
  66. def main():
  67.     loop = asyncio.get_event_loop()
  68.  
  69.     listen = loop.create_datagram_endpoint(UDPServerProtocol, LISTEN)
  70.     udp, proto = loop.run_until_complete(listen)
  71.  
  72.     listen = asyncio.start_server(tcp_forwarder, *LISTEN)
  73.     tcp = loop.run_until_complete(listen)
  74.  
  75.     print('running')
  76.     try:
  77.         loop.run_forever()
  78.     except KeyboardInterrupt:
  79.         pass
  80.     udp.close()
  81.     tcp.close()
  82.     loop.close()
  83.  
  84. if __name__ == '__main__':
  85.     main()
RAW Paste Data
Top