#!/bin/python
from dnslib import server,RR,QTYPE
#import dnslib
import re
import dns.query
import dns.message
import netifaces
import os
import argparse
class bcolors:
    OK = '\033[92m' #GREEN
    WARNING = '\033[93m' #YELLOW
    FAIL = '\033[91m' #RED
    RESET = '\033[0m' #RESET COLOR

parser = argparse.ArgumentParser(description='a DoT Proxy resolver.')
parser.add_argument('-r',action="store_true", help='Optional: Enable request intercepting')
args = parser.parse_args()
if args.r == True:
    print(bcolors.OK+"Domain intercepting enabled!"+bcolors.RESET)
listeningIPs = []
print("DoT (DNS over TLS) Proxy resolver - <Written by Caleb Fontenot>")
# ASCII Art because why not
operating_system = os.name
print(operating_system)
if operating_system == "posix":
    os.system("cat ascii")
elif operating_system == "nt":
    os.system("type ascii")
#Get IPs of interfaces
for iface in netifaces.interfaces():
  iface_details = netifaces.ifaddresses(iface)
  if netifaces.AF_INET in iface_details:
    interface_info1 = iface_details[netifaces.AF_INET]
    interface_info = interface_info1[0]
    listeningIPs.append(interface_info["addr"])
print("Listening on", listeningIPs)
print("Hailey is my Soulmate!")
ttl=53

# 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')'
IPV6SEG  = r'(?:(?:[0-9a-fA-F]){1,4})'
IPV6GROUPS = (
    r'(?:' + IPV6SEG + r':){7,7}' + IPV6SEG,                  # 1:2:3:4:5:6:7:8
    r'(?:' + IPV6SEG + r':){1,7}:',                           # 1::                                 1:2:3:4:5:6:7::
    r'(?:' + IPV6SEG + r':){1,6}:' + IPV6SEG,                 # 1::8               1:2:3:4:5:6::8   1:2:3:4:5:6::8
    r'(?:' + IPV6SEG + r':){1,5}(?::' + IPV6SEG + r'){1,2}',  # 1::7:8             1:2:3:4:5::7:8   1:2:3:4:5::8
    r'(?:' + IPV6SEG + r':){1,4}(?::' + IPV6SEG + r'){1,3}',  # 1::6:7:8           1:2:3:4::6:7:8   1:2:3:4::8
    r'(?:' + IPV6SEG + r':){1,3}(?::' + IPV6SEG + r'){1,4}',  # 1::5:6:7:8         1:2:3::5:6:7:8   1:2:3::8
    r'(?:' + IPV6SEG + r':){1,2}(?::' + IPV6SEG + r'){1,5}',  # 1::4:5:6:7:8       1:2::4:5:6:7:8   1:2::8
    IPV6SEG + r':(?:(?::' + IPV6SEG + r'){1,6})',             # 1::3:4:5:6:7:8     1::3:4:5:6:7:8   1::8
    r':(?:(?::' + IPV6SEG + r'){1,7}|:)',                     # ::2:3:4:5:6:7:8    ::2:3:4:5:6:7:8  ::8       ::
    r'fe80:(?::' + IPV6SEG + r'){0,4}%[0-9a-zA-Z]{1,}',       # fe80::7:8%eth0     fe80::7:8%1  (link-local IPv6 addresses with zone index)
    r'::(?:ffff(?::0{1,4}){0,1}:){0,1}[^\s:]' + IPV4ADDR,     # ::255.255.255.255  ::ffff:255.255.255.255  ::ffff:0:255.255.255.255 (IPv4-mapped IPv6 addresses and IPv4-translated addresses)
    r'(?:' + IPV6SEG + r':){1,4}:[^\s:]' + IPV4ADDR,          # 2001:db8:3:4::192.0.2.33  64:ff9b::192.0.2.33 (IPv4-Embedded IPv6 Address)
)
IPV6ADDR = '|'.join(['(?:{})'.format(g) for g in IPV6GROUPS[::-1]])  # Reverse rows for greedy match
regex4 = re.compile(IPV4ADDR)
regex6 = re.compile(IPV6ADDR)
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])        
        for i in range(len(IP)):
            query_response = str(request.q.qname)+" "+str(ttl)+" "+str(QTYPE[request.q.qtype])+" "+IP[i]
            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)
try:
    server = server.DNSServer(resolver,port=53,address="0.0.0.0",logger=logger, tcp=False)
except:
    print(bcolors.FAIL+"Failed to open socket!"+bcolors.RESET)
    if operating_system == "posix":
        print("Detected posix environment. Linux requires root privleges to open port 53")
        exit()
#server.start_thread()
    
# Query upstream DoH server
def DoTquery(domain, query_type):
    if args.r == True:
        if "nintendowifi.net." in str(domain):
            print(bcolors.FAIL+"Intercepted query!"+bcolors.RESET)
            return ["172.104.88.237"]
    q = dns.message.make_query(str(domain), str(query_type))
    response = dns.query.tls(q, "76.72.69.144", server_hostname="resolve.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":    
        IP = regex4.findall(str(response.answer))
        print(bcolors.OK+"IP is", str(IP[0])+bcolors.RESET)
    elif str(query_type) == "AAAA":
        IP = regex6.findall(str(response.answer))
        #print("Not implemented yet!")
    else:
        print(bcolors.FAIL+"Queries beyond A or AAAA are not implemented"+bcolors.RESET)
    return IP
#server.start_thread()
server.start()