Skip to content

Commit

Permalink
Standardized Queries
Browse files Browse the repository at this point in the history 
Standardized the queried, still need to add proper error checking on the queries and return SERVFAIL
  • Loading branch information
@mpretzel16
mpretzel16 committed Oct 3, 2021
1 parent e272533 commit 219e1bf
Show file tree
Hide file tree
Showing 9 changed files with 321 additions and 312 deletions.
184 changes: 28 additions & 156 deletions dns_server.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,12 +23,12 @@ class DNSServer:
conn = None
arr_processing = []
a_record: A
aaaa_record = AAAA()
mx_record = MX()
txt_record = TXT()
cname_record = CNAME()
ptr_record = PTR()
ns_record = NS()
aaaa_record: AAAA
mx_record: MX
txt_record: TXT
cname_record: CNAME
ptr_record: PTR
ns_record: NS
database: Database
dict_server_config: dict
def __init__(self):
Expand All @@ -42,6 +42,12 @@ def __init__(self):
exit(1)
self.dict_server_config = config['dnsServer']
self.a_record = A(self.dict_server_config)
self.aaaa_record = AAAA(self.dict_server_config)
self.ns_record = NS(self.dict_server_config)
self.mx_record = MX(self.dict_server_config)
self.cname_record = CNAME(self.dict_server_config)
self.txt_record = TXT(self.dict_server_config)
self.ptr_record = PTR(self.dict_server_config)

def handle_request(self, request, address):
start_time = time.time()
Expand All @@ -51,161 +57,27 @@ def handle_request(self, request, address):
self.a_record.handle_request(dns, self.database, address, self.sock)
elif dns[DNSQR].qtype == 2: # NS Record
self.ns_record.handle_request(dns, self.database, address, self.sock)
elif dns[DNSQR].qtype == 5:
self.cname_record.handle_request(dns, self.conn)
elif dns[DNSQR].qtype == 12:
elif dns[DNSQR].qtype == 5: # CNAME Record
self.cname_record.handle_request(dns, self.database, address, self.sock)
elif dns[DNSQR].qtype == 12: # PTR Record
self.ptr_record.handle_request(dns, self.database, address, self.sock)
elif dns[DNSQR].qtype == 15:
self.mx_record.handle_request(dns, self.conn)
elif dns[DNSQR].qtype == 16:
self.txt_record.handle_request(dns, self.conn, address, self.sock)
elif dns[DNSQR].qtype == 15: # MX Record
self.mx_record.handle_request(dns, self.database, address, self.sock)
elif dns[DNSQR].qtype == 16: # TXT Record
self.txt_record.handle_request(dns, self.database, address, self.sock)
elif dns[DNSQR].qtype == 28: # AAAA Record
self.aaaa_record.handle_request(dns, self.database, address, self.sock)


# # Define Postgres Functions
# postgres = PostgresFunctions()
# # Set DNS Request to a variable
# dns = DNS(request)
# # Get the ascii string for the DNS request
# query = str(dns[DNSQR].qname.decode('ascii'))
# if query.lower().endswith('.pretzel1.'):
# response = DNS(id=dns.id, ancount=1, qr=1, rcode=3)
# ans4 = DNSRR(rrname=str(query), type=1)
# response.an = ans4
# self.sock.sendto(bytes(response), addr)
# return True
# # Check if this server is Authoritative for this DNS request (It will be 99% of the time)
# result_dns_authority = postgres.get_dns_authority(self.conn, str(query))
# # Error Checking Authoritative status, do not send a result
# if result_dns_authority['bool_error']:
# return False
# # Put database results into a variable
# arr_data = json.loads(result_dns_authority['data'])
# # No data was returned from the SQL query (Non Existent Domain)
# if len(arr_data) == 0:
# # We will only auto create records for A and AAA types
# if dns[DNSQR].qtype == 1 or dns[DNSQR].qtype == 28:
# # Generate IPs for A and AAA, add to database
# dict_result_generate_dns_entry = GenerateDnsEntry().start(self.conn, str(query))
# # Error in creating new DNS record, do not send result
# if dict_result_generate_dns_entry['bool_error']:
# return False
# # All set to return generated data
# if not dict_result_generate_dns_entry['bool_already_exists']:
# response = DNS(id=dns.id, ancount=1, qr=1)
# ans4 = DNSRR(rrname=str(query), type=1, rdata=str(dict_result_generate_dns_entry['ipv4']))
# response.an = ans4
# self.sock.sendto(bytes(response), addr)
# return True
# else: # A record was created before this record could be created. Reference new record
# if dns[DNSQR].qtype == 1:
# result_dns_entry = postgres.get_dns_record(self.conn, str(query), 1)
# result_data = json.loads(result_dns_entry['data'])
# int_total = len(result_data)
# response = DNS(id=dns.id, ancount=int(int_total), qr=1)
# if int_total > 0:
# int_parsed = 0
# for result in result_data:
# #print(result)
# ans = DNSRR(rrname=str(query), type=1, rdata=str(result['record_data']['ip']))
# if int_parsed == 0:
# response.an = ans
# elif int_parsed < int_total:
# response.an = response.an / ans
# int_parsed += 1
# else:
# ans = DNSRR(rrname=str(query), type=1, rdata=None)
# response.an = ans
# self.sock.sendto(bytes(response), addr)
# return True
# if dns[DNSQR].qtype == 28:
# result_dns_entry = postgres.get_dns_record(self.conn, str(query), 28)
# result_data = json.loads(result_dns_entry['data'])
# int_total = len(result_data)
# response = DNS(id=dns.id, ancount=int(int_total), qr=1)
# if int_total > 0:
# int_parsed = 0
# for result in result_data:
# #(result['record_data']['ip'])
# ans = DNSRR(rrname=str(query), type=28, rdata=str(result['record_data']['ip']))
# if int_parsed == 0:
# response.an = ans
# elif int_parsed < int_total:
# response.an = response.an / ans
# int_parsed += 1
# else:
# ans = DNSRR(rrname=str(query), type=1, rdata=None)
# response.an = ans
# self.sock.sendto(bytes(response), addr)
# return True
# if arr_data[0]['bool_authoritative'] == "1":
# if dns[DNSQR].qtype == 1:
# result_dns_entry = postgres.get_dns_record(self.conn, str(query), 1)
# result_data = json.loads(result_dns_entry['data'])
# int_total = len(result_data)
# response = DNS(id=dns.id, ancount=int(int_total), qr=1)
# if int_total > 0:
# int_parsed = 0
# for result in result_data:
# #print(result)
# ans = DNSRR(rrname=str(query), type=1, rdata=str(result['record_data']['ip']))
# if int_parsed == 0:
# response.an = ans
# elif int_parsed < int_total:
# response.an = response.an / ans
# int_parsed += 1
# else:
# ans = DNSRR(rrname=str(query), type=1, rdata=None)
# response.an = ans
# self.sock.sendto(bytes(response), addr)
# if dns[DNSQR].qtype == 28:
# result_dns_entry = postgres.get_dns_record(self.conn, str(query), 28)
# result_data = json.loads(result_dns_entry['data'])
# int_total = len(result_data)
# response = DNS(id=dns.id, ancount=int(int_total), qr=1)
# if int_total > 0:
# int_parsed = 0
# for result in result_data:
# #print(result['record_data']['ip'])
# ans = DNSRR(rrname=str(query), type=28, rdata=str(result['record_data']['ip']))
# if int_parsed == 0:
# response.an = ans
# elif int_parsed < int_total:
# response.an = response.an / ans
# int_parsed += 1
# else:
# ans = DNSRR(rrname=str(query), type=1, rdata=None)
# response.an = ans
# self.sock.sendto(bytes(response), addr)
# if dns[DNSQR].qtype == 12:
# result_dns_entry = postgres.get_dns_record(self.conn, str(query), 12)
# result_data = json.loads(result_dns_entry['data'])
# int_total = len(result_data)
# response = DNS(id=dns.id, ancount=int(int_total), qr=1)
# #print(int_total)
# if int_total > 0:
# int_parsed = 0
# for result in result_data:
# ans = DNSRR(rrname=str(query), type=12, rdata=str(result['record_data']['domain_name']))
# if int_parsed == 0:
# response.an = ans
# elif int_parsed < int_total:
# response.an = response.an / ans
# int_parsed += 1
# else:
# ans = DNSRR(rrname=str(query), type=1, rdata=None)
# response.an = ans
# self.sock.sendto(bytes(response), addr)
# #print("--- %s seconds ---" % (time.time() - start_time))
# return False
# #domain, tld, tail = query.rsplit('.', 3)
# #assert domain == 'example' and tld == 'com' and tail == ''
# #head = head.split('.', 1)[-1] # drop leading "prefix." part
elif dns[DNSQR].qtype == 255: # ANY Record
pass
#self.aaaa_record.handle_request(dns, self.database, address, self.sock)
else:
print(dns[DNSQR].qtype)
response = DNS(id=dns.id, ancount=0, qr=1)
self.sock.sendto(bytes(response), address)
except Exception as e:
print('')
print(e)
print_exc()
print('garbage from {!r}? data {!r}'.format(addr, request))
print('garbage from {!r}? data {!r}'.format(address, request))


def run_server(self):
Expand Down
50 changes: 20 additions & 30 deletions generate_dns_entry.py
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,7 @@
import psycopg2
import sqlalchemy.orm

from random import randint
import json
from ip_generation import IpGeneration
from sqlalchemy import func
from sqlalchemy import or_
from database import Database


Expand All @@ -29,33 +25,11 @@ def start(self, db_session: sqlalchemy.orm.Session, entry_text):
dict_return['bool_error'] = True
dict_return['str_status'] = dict_tld_query['str_status']
return dict_return
lookup_tld = db_session.query(func.count(Database.Tld_Listing.domain))\
lookup_tld = db_session.query(func.count(Database.Tld_Listing.domain)) \
.filter(Database.Tld_Listing.domain.in_(dict_tld_query['arr_tld'])).scalar()
if lookup_tld == 0:
dict_return['bool_invalid_tld'] = True
return dict_return
try:
dns_authority = Database.Dns_Authoritative()
dns_authority.bool_authoritative = True
dns_authority.entry_text = entry_text
db_session.add(dns_authority)
db_session.commit()
except sqlalchemy.exc.IntegrityError as exec:
if "duplicate" in str(exec).lower():
dict_return['bool_already_exists'] = True
return dict_return
else:
dict_return['bool_error'] = True
dict_return['str_status'] = str(exec)
return dict_return
except Exception as e:
if "duplicate" in str(e).lower():
dict_return['bool_already_exists'] = True
return dict_return
dict_return['bool_error'] = True
dict_return['str_status'] = str(e)
return dict_return

country_to_use: Database.Country_List
try:
result_country_list = db_session.query(Database.Country_List).all()
Expand All @@ -70,7 +44,7 @@ def start(self, db_session: sqlalchemy.orm.Session, entry_text):
return dict_return
result_ipv4: Database.IPv4_Networks
try:
result_ipv4 = db_session.query(Database.IPv4_Networks)\
result_ipv4 = db_session.query(Database.IPv4_Networks) \
.filter(Database.IPv4_Networks.country_id == country_to_use.country_id).all()
except Exception as e:
dict_return['bool_error'] = True
Expand Down Expand Up @@ -112,12 +86,28 @@ def start(self, db_session: sqlalchemy.orm.Session, entry_text):
ns_dns_entry = Database.Dns_Entries()
ns_dns_entry.entry_text = entry_text
ns_dns_entry.record_type = 2
ns_dns_entry.record_data = {"server_id": self.dict_server_config['id']}
db_session.bulk_save_objects([a_dns_entry, aaa_dns_entry, ptr_a_dns_entry, ptr_aaa_dns_entry, ns_dns_entry])
ns_dns_entry.record_data = {"server_id": self.dict_server_config['id'],
"domain_name": self.dict_server_config['hostname']}
ns_a_dns_entry = Database.Dns_Entries()
ns_a_dns_entry.entry_text = str_ptr_ipv4
ns_a_dns_entry.record_type = 2
ns_a_dns_entry.record_data = {"server_id": self.dict_server_config['id'],
"domain_name": self.dict_server_config['hostname']}
ns_aaaa_dns_entry = Database.Dns_Entries()
ns_aaaa_dns_entry.entry_text = str_ptr_ipv6
ns_aaaa_dns_entry.record_type = 2
ns_aaaa_dns_entry.record_data = {"server_id": self.dict_server_config['id'],
"domain_name": self.dict_server_config['hostname']}
db_session.bulk_save_objects([a_dns_entry, aaa_dns_entry, ptr_a_dns_entry, ptr_aaa_dns_entry, ns_dns_entry,
ns_a_dns_entry, ns_aaaa_dns_entry])
db_session.commit()
dict_return['ipv4'] = ips['ipv4']
dict_return['ipv6'] = ips['ipv6']
except Exception as e:
print(e)
if "duplicate" in str(e).lower():
dict_return['bool_already_exists'] = True
return dict_return
dict_return['bool_error'] = True
dict_return['str_status'] = str(e)
return dict_return
Expand Down
Loading

0 comments on commit 219e1bf

Please sign in to comment.