#!/bin/python
from dnslib import server,RR,QTYPE
#import dnslib
import re
import dns.query
import dns.message
print("DoT (DNS over TLS) Proxy resolver - <Written by Caleb Fontenot>")
ttl=53
class bcolors:
    OK = '\033[92m' #GREEN
    WARNING = '\033[93m' #YELLOW
    FAIL = '\033[91m' #RED
    RESET = '\033[0m' #RESET COLOR

# REGEX Magic, how does it work
IPV4SEG  = r'(?:25[0-5]|(?:2[0-4]|1{0,1}[0-9]){0,1}[0-9])'
IPV4ADDR = r'(?:(?:' + IPV4SEG + r'\.){3,3}' + IPV4SEG + r')'
class DoTProxy:
    def resolve(self,request,handler):
        reply = request.reply()
        print(bcolors.WARNING, request.q.qname, str(QTYPE[request.q.qtype]), bcolors.RESET)
        IP = DoTquery(str(request.q.qname), QTYPE[request.q.qtype])
        query_response = str(request.q.qname)+" "+str(ttl)+" "+str(QTYPE[request.q.qtype])+" "+IP
        print(str(query_response))
        reply.add_answer(*RR.fromZone(query_response))
        return reply
    
#Setup DNS server, this will listen for incoming DNS packets
resolver = DoTProxy()
logger = server.DNSLogger(prefix=False)
server = server.DNSServer(resolver,port=53,address="0.0.0.0",logger=logger, tcp=False)
#server.start_thread()
# Query upstream DoH server
def DoTquery(domain, query_type):
    #response = pythonDoh.client.query(str(domain), type=str(query_type), server="calebfontenot.com", verbose=True, fallback=False)
    if str(domain) == "conntest.nintendowifi.net.":
        print("bruh moment")
        return "172.104.88.237"
    else:
        q = dns.message.make_query(str(domain), str(query_type))
        response = dns.query.tls(q, "74.80.18.217", server_hostname="calebfontenot.com", port=853)
        print(response)
        # There's probably a better way of doing this, but I can't figure it out...
        if str(query_type) == "A":    
            regex = re.compile(IPV4ADDR)
            IP = regex.findall(str(response.answer))
            print(bcolors.OK+"IP is", str(IP[0])+bcolors.RESET)
        elif str(query_type) == "AAAA":
            print("Not implemented yet!")
        else:
            return("Queries beyond A or AAAA are not implemented")
        return IP[0]
#server.start_thread()
server.start()