#!/usr/bin/env python3
# 20181111
# Jan Mojzis
# Public domain.

# Simple DNS responder
# For TXT query sends large response (~64kB)
# For other queries sends small A response

import logging
import socket
import sys
import select
import struct

T_TXT = b'\x00\x10'
T_A = b'\x00\x01'
C_IN = b'\x00\x01'

def be16(x = 0):
        """
        pack uint16_t to big-endian
        """

        return struct.pack('>H', x)

def be32(x = 0):
        """
        pack uint32_t to big-endian
        """

        return struct.pack('>I', x)

def recvtcp(c):
        """
        """
        ret = b''

        data = c.recv(2)
        l = struct.unpack('>H', data)[0]

        while len(ret) != l:
                data = c.recv(65536)
                ret += data
        return ret

def sendtcp(c, response):
        """
        """

        l = len(response)
        c.sendall(be16(l) + response)


if __name__ == '__main__':
        logging.basicConfig(level=logging.DEBUG)

        try:
                ip = sys.argv[1]
        except:
                ip = '127.0.0.2'

        try:
                port = int(sys.argv[2])
        except:
                port = 53

        logging.info("Starting fakednserver: %s:%d", ip, port)

        udp = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        udp.bind((ip, port))

        tcp = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        tcp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        tcp.bind((ip, port))
        tcp.listen(10)

        while True:
                r, w, e = select.select([udp, tcp], [], [], None)
                for s in r:
                        t = None
                        if s == tcp:
                                c, addr = tcp.accept()
                                data = recvtcp(c)
                                t = 'tcp'
                        elif s == udp:
                                data, addr = udp.recvfrom(4096)
                                t = 'udp'
                        else:
                                continue

                        domain = b''
                        opcode = (data[2] >> 3) & 15
                        if opcode != 0:
                                continue

                        start = 12
                        length = data[start]
                        while length != 0:
                                domain += data[start + 1:start + length + 1] + b'.'
                                start += length + 1
                                length = data[start]
                        qtype = data[start + 1: start + 3]

                        #XXX drop tail
                        data = data[0:start + 5]

                        logging.info("DNS %s request from: %s: %s", t, addr, domain)

                        if qtype == T_TXT:
                                answers = 240
                        else:
                                answers = 1

                        response = b''
                        response += data[:2]            # ID
                        response += b"\x81\x80"         # FLAGS
                        response += data[4:6]           # QUESTION
                        response += be16(answers)       # ANSWER
                        response += be16(0)             # AUTHORITY
                        response += be16(0)             # ADDITIONAL
                        response += data[12:]           # NAME

                        truncated = b''
                        truncated += data[:2]
                        truncated += b"\x83\x80"
                        truncated += response[4:]

                        if qtype == T_TXT:
                                # long TXT response over TCP
                                for i in range(answers):
                                        response += b'\xc0\x0c'         # NAME (pointer to name)
                                        response += T_TXT               # TYPE TXT
                                        response += C_IN                # CLASS IN
                                        response += be32(60)            # TTL
                                        txtlen = 255
                                        response += be16(txtlen)        # RDLENGTH
                                        response += txtlen * b'\xfe'    # RDATA (note: first byte is the length)

                        else:
                                # short A response over UDP
                                for i in range(answers):
                                        response += b'\xc0\x0c'         # NAME (pointer to name)
                                        response += T_A                 # TYPE A
                                        response += C_IN                # CLASS IN
                                        response += be32(60)            # TTL
                                        response += be16(4)             # RDLENGTH
                                        response += b'\x00\x00\x00\x00' # RDATA

                        if s == tcp:
                                sendtcp(c, response)
                                c.close()
                        if s == udp:
                                if len(response) > 512:
                                        udp.sendto(truncated, addr)
                                else:
                                        udp.sendto(response, addr)
