#!/usr/bin/env python

"""
udpproxy - Tool to handle the UDP side of DNS communication

This program is meant to accompany the PROTOS DNS test material.
Run it on the same computer that the DNS jar is running on, and
it will handle UDP->TCP translation of DNS messages so that
the test tool can receive DNS messages sent over UDP.

Command line options:

-h  --help               This text
-l  --listenport NNN     UDP port to listen on
-e  --testtoolport NNN   TCP port used by the test tool on this computer
-a  --targetport  NNN    TCP port of target, for TCP replies
-d  --duptimeout         Discard duplicate packets received within 1 second
-x  --hexdumps           Display hex dumps of proxied traffic
                         (requires the "xxd" program in path)

All ports default to 53 (standard DNS port), and the default will work unless
you have configured  the test tool or test subject to use nonstandard port(s).

This program can be distributed under the GNU LGPL
(http://www.gnu.org/licenses/lgpl.html).

"""

from socket import *
from socket import error as SocketError
import sys, struct
import errno, time
import datetime
import getopt
import string
import os

def timestamp():
    return datetime.datetime.now().isoformat()

def log(*args):
    args = map(str, args)
    
    print timestamp()+' ' + string.join(args)

def hexdump(data):
    p = os.popen('xxd', 'w')
    p.write(data)
    p.close()

def recv_n(s, n):
    out = ''
    while len(out) < n:
        x = s.recv(n-len(out))
        if not x:
            break
        else:
            out += x
    if len(out) != n:
        if len(out) > 0:
            log("did not get %d bytes from socket (got %d)" % (n, len(out)))
        raise EOFError()
    return out

def read_one(s):
    reqlen = struct.unpack('!H', recv_n(s, 2))[0]
    if reqlen == 0:
        # log("zero length dns packet from tcp socket, skip...")
        return ''
    log("reading dns packet of len %s from tcp socket" % reqlen)
    potk = recv_n(s, reqlen)
    return potk

def tcpquery(data):
    ntries = 10
    for i in range(ntries):
        s = socket(AF_INET, SOCK_STREAM)
        try:
            s.connect(('localhost', testtoolport))
        except SocketError, what:
            (nr, msg) = what
            if nr == errno.ECONNREFUSED:
                # log('connection refused on localhost:%d. sleeping before retry...' % testtoolport)
                time.sleep(1)
                continue
            else:
                log("socket error %s writing dns message to test tool" % str(what))

        try:
            s.sendall(struct.pack('!H', len(data)) + data)
        except SocketError, what:
            (nr, msg) = what
            log("got %s sending data to test tool, stop" % str(what))
            return []

        answers = []
        while 1:
            try:
                x = read_one(s)
            except EOFError:
                log("end-of-file receiving tcp message from test tool")
                break
            except SocketError, what:
                log("tcp read error from test tool socket: %s" % what)
                break
            else:
                answers.append(x)
        s.close()
        return answers
    log('give up after %s tries -- cannot connect to test tool' % ntries)
    return []

def tcpreply(data, a):
    if not data:
        log("empty message, not sending")
        return
    s = socket(AF_INET, SOCK_STREAM)
    log('send', len(data), 'to', a[0], '...')
    try:
        s.connect((a[0], targetport))
    except SocketError, (nr, msg):
        if nr != errno.ECONNREFUSED:
            log("can't open tcp connection to deliver reply: %s" % str(msg))
            return
        else:
            log("sending via udp since tcp connection was refused")
            try:
                udps.sendto(data, a)
            except SocketError, (nr, msg):
                log("udp send failed: %s" % str(msg))
            return
    try:
        s.sendall(struct.pack('!H', len(data)) + data)
        s.close()
    except SocketError, what:
        log("tcp write error: %s" % what)
    log('done')
    sys.stdout.flush()

try:
    opts, rest = getopt.getopt(sys.argv[1:], 'hl:a:e:dx',
                               ['help', 'listenport=', 'testtoolport=',
                                'targetport=', 'duptimeout', 'hexdumps'])
except getopt.error, what:
    sys.exit(str(what))

listenport = testtoolport = targetport = 53
duptimeout = 0
wanthexdumps = 0
for o, a in opts:
    if o in ('-h', '--help'):
        print __doc__
        sys.exit()
    elif o in ('-l', '--listenport'):
        listenport = int(a)
    elif o in ('-e', '--testtoolport'):
        testtoolport = int(a)
    elif o in ('-a', '--targetport'):
        targetport = int(a)
    elif o in ('-d', '--duptimeout'):
        duptimeout=1
    elif o in ('-x', '--hexdumps'):
        wanthexdumps=1

prevdata=None
prevtime = time.time()
udps = socket(AF_INET, SOCK_DGRAM)
udps.bind(('', listenport))


while 1:
    try:
        data, srcaddr = udps.recvfrom(1500)
    except SocketError, what:
        log("udp receive failed: %s" % what)
        continue

    log('received udp request of length %s from %s:%s' % (len(data), srcaddr[0], srcaddr[1]))
    if wanthexdumps:
        hexdump(data)
    if data == prevdata:
        if duptimeout and ((time.time()-prevtime) > 1):
            continue
    try:
        answers = tcpquery(data)
        log('did tcpquery')
    except SocketError, what:
        (nr, msg) = what
        log(msg)
        log("traceback:")
        import traceback
        traceback.print_exc()
    else:
        for a in answers:
            if not a:
                continue
            try:
                tcpreply(a, srcaddr)
                log('sent tcp reply of length', len(a))
                if wanthexdumps:
                    hexdump(a)
            except SocketError, what:
                log("socket error while sending reply: %s" % what)
        prevdata= data
        prevtime = time.time()