#!/usr/bin/env python3

import socket
import logging
import select
import struct
from dnslib import *

def rootserver(query):
    """
    fake root server
    """

    #logging.debug('%s' % (query.header))
    #logging.debug('%s' % (query.question))
    #logging.debug('%s' % (query.q.qname))
    #logging.debug('%s' % (nsname))

    qname = query.q.qname
    nsname = 'ns.' + str(qname).split('.')[-2]

    answer = query.reply()
    nsip='127.0.0.3'
    answer.add_auth(RR(qname,QTYPE.NS,rdata=NS(nsname),ttl=3600))
    answer.add_ar(RR(nsname,QTYPE.A,rdata=A(nsip),ttl=3600))

    logging.debug('query: %s, UDP: auth(%s), ar(%s)' % (qname, nsname, nsip))
    return answer

def truncserver(query):
    """
    fake DNS server, sends truncated answer
    """

    qname = query.q.qname
    answer = query.reply()
    ip='127.0.0.3'
    logging.debug('query: %s, UDP: truncated' % (qname))
    answer.header.tc = 1
    return answer

def server(query):
    """
    fake DNS server, sends A response
    """

    qname = query.q.qname
    answer = query.reply()
    ip='127.0.0.3'
    answer.add_answer(RR(qname,QTYPE.A,rdata=A(ip),ttl=3600))
    answer.add_ar(RR('dummy',QTYPE.A,rdata=A(ip),ttl=3600))
    logging.debug('query: %s, TCP: answer(%s)' % (qname, ip))
    return answer


if __name__ == '__main__':

    logging.basicConfig(level=logging.DEBUG)

    # create socket for fake root DNS server
    ip = '127.0.0.2'
    port = 53
    rootudp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    rootudp.bind((ip, port))
    logging.info(f'running fake root DNS server %s:%s' % (ip, port))

    # create socket for fake DNS server
    ip = '127.0.0.3'
    port = 53
    nsudp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    nsudp.bind((ip, port))
    logging.info(f'running fake DNS server %s:%s' % (ip, port))

    # create socket for fake DNS server
    ip = '127.0.0.3'
    port = 53
    nstcp = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    nstcp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    nstcp.bind((ip, port))
    nstcp.listen()

    while True:
        # read DNS query
        try:
            (rr, ww, ee) = select.select([rootudp, nsudp, nstcp], [], [])
            for r in rr:
                if r in [rootudp, nsudp]:
                    packet, addr = r.recvfrom(512)
                else:
                    c, addr = r.accept()
                    # read the length
                    packet = c.recv(1)
                    packet += c.recv(1)
                    length = struct.unpack("!H",bytes(packet))[0]
                    # read the packet
                    packet = b''
                    for i in range(length):
                        packet += c.recv(1)

                # parse DNS query
                query = DNSRecord.parse(packet)
                qname = query.q.qname

                if r == rootudp:
                    answer = rootserver(query)
                elif r == nsudp:
                    answer = truncserver(query)
                elif r == nstcp:
                    answer = server(query)
                else:
                    raise("XXX")

                # send response to the client
                if r in [rootudp, nsudp]:
                    r.sendto(answer.pack(), addr)
                else:
                    data = answer.pack()
                    packet = struct.pack("!H",len(data)) + data
                    c.sendall(packet, len(packet))
                    c.close()
        except KeyboardInterrupt:
            break
        except Exception:
            pass
