#!/usr/bin/env /usr/bin/python
# -*- coding: utf-8 -*-

"""
  FileResolver - example resolver which responds with fixed response
          to all requests
"""

import os
import os.path
import base64

from dnslib import RR, QTYPE, TXT
from dnslib.server import DNSServer, DNSHandler, BaseResolver, DNSLogger


class FileResolver(BaseResolver):
    """
      Respond with fixed response to all requests
    """

    def __init__(self, directory, domain):
        self.filelist = []
        self.domain = domain
        self.ttl = 60
        self.directory = directory
        self.cache = {}
        if directory:
            if not os.path.isdir(directory):
                print(f"Directory {directory} doesn't exist")
                exit()

    def getcache(self, path):
        pname = os.path.basename(path)
        if pname in self.cache:
            # it's in the cache - so lets just read it from the cache
            print("Taking from cache")
            cooked = self.cache[pname]["base64"]
            # And update time
            self.cache[pname]["time"] = time.time()
        else:
            # It's not in cache - read the file and cache it
            print("Adding to cache")
            fin = open(path, "rb")
            raw = fin.read()
            fin.close()
            cooked = base64.b64encode(raw)
            self.cache[pname] = {}
            self.cache[pname]["base64"] = cooked
            self.cache[pname]["time"] = time.time()

        self.cachecheck()
        return cooked

    def cachecheck(self):
        # Remove any entries in the cache over an hour old
        curtime = time.time()
        staletime = 3600
        for key in self.cache:
            if curtime - self.cache[key]["time"] > staletime:
                # it's stale
                del self.cache[key]

    def resolve(self, request, handler):
        qry_type = QTYPE[request.q.qtype]
        name = request.q.qname
        reply = request.reply()

        # First check domain is at the end of the name
        if not name.matchSuffix(self.domain):
            reply.add_answer(RR(name, QTYPE.TXT, ttl=self.ttl, rdata=TXT("Domain not found")))
            return reply

        # Request method for downloads
        if qry_type == "TXT":
            # Format is filename.count.domain for count
            if not name.matchSuffix(self.domain):
                reply.add_answer(RR(name, QTYPE.TXT, ttl=self.ttl, rdata=TXT("Domain not found")))
                return reply

            # Now look whether we're looking for a count or a part
            parts = str(name.stripSuffix(self.domain)).split(".")
            # parts[0] should be filename, parts[1] should be segment or count
            pname = '.'.join(parts[:-2])
            path = self.directory + "/" + pname
            command = ''.join(parts[-2:-1])
            # TODO: circumvent caching, eg by counting up
            if command.endswith("u"):
                with open('/tmp/dns-srv.log', 'a') as f:
                    f.write(str(pname) + "\n")
                reply.add_answer(RR(name, QTYPE.TXT, ttl=self.ttl, rdata=TXT("Upload accepted")))
                return reply
            if not os.path.isfile(path):
                reply.add_answer(RR(name, QTYPE.TXT, ttl=self.ttl, rdata=TXT("File not found")))
                return reply

            # work out the count mathematically
            # First work out the base64 size
            length = os.path.getsize(path)
            length = (4 * length) / 3
            # And padding
            length += (length % 4)
            # Finally divide into number of 254 byte chunks
            chunks = (length / 254)

            if command == "count":
                reply.add_answer(RR(name, QTYPE.TXT, ttl=self.ttl, rdata=TXT(str(chunks))))
                return reply

            if command.isdigit():
                # Woo it's a number
                # lets base64 the file
                chunk = int(command)
                if chunk > chunks or chunk < 0:
                    reply.add_answer(RR(name, QTYPE.TXT, ttl=self.ttl, rdata=TXT("Chunk out of range")))
                    return reply
                # It's actually trying to read the file, so let's cache it in memory
                cooked = self.getcache(path)

                # Now lets just grab the chunk
                start = chunk * 254
                txtr = cooked[start:start + 254]
                reply.add_answer(RR(name, QTYPE.TXT, ttl=self.ttl, rdata=TXT(txtr)))
                return reply

        # Replace labels with request label
        return reply


if __name__ == '__main__':
    import argparse
    import time

    p = argparse.ArgumentParser(description="Fixed DNS Resolver")
    p.add_argument("--directory", "-r", metavar="<directory>", help="Directory to pull files from")
    p.add_argument("--domain", "-d", metavar="<domain>", help="Domain we're working in")
    p.add_argument("--port", "-p", type=int, default=53, metavar="<port>", help="Server port (default:53)")
    p.add_argument("--address", "-a", default="", metavar="<address>", help="Listen address (default:all)")
    p.add_argument("--udplen", "-u", type=int, default=0, metavar="<udplen>", help="Max UDP packet length (default:0)")
    p.add_argument("--tcp", action='store_true', default=False, help="TCP server (default: UDP only)")
    p.add_argument("--log", default="request,reply,truncated,error",
                   help="Log hooks (default: +request,+reply,+truncated,+error,-recv,-send,-data)")
    p.add_argument("--log-prefix", action='store_true', default=False, help="Log prefix (default: False)")
    args = p.parse_args()

    resolver = FileResolver(args.directory, args.domain)
    logger = DNSLogger(args.log, args.log_prefix)

    print(f"Starting File Resolver ({args.address or '*'}:{args.port}) [{'UDP/TCP' if args.tcp else 'UDP'}]")

    if args.udplen:
        DNSHandler.udplen = args.udplen

    udp_server = DNSServer(resolver, port=args.port, address=args.address, logger=logger)
    udp_server.start_thread()

    if args.tcp:
        tcp_server = DNSServer(resolver, port=args.port, address=args.address, tcp=True, logger=logger)
        tcp_server.start_thread()

    while udp_server.isAlive():
        time.sleep(1)