from scapy.all import *

''' 
	Make DNS query to Google at 8.8.8.8 using Python

	Author: Tim Pierson, Dartmouth CS60, Fall 2026

	run: sudo python3 google_dns.py <domain name> [A|MX|NS]
			sudo python3 google_dns.py dartmouth.edu.      --> returns dartmouth's A record
			sudo python3 google_dns.py dartmouth.edu. NS   --> returns dartmouth's nameservers
'''

def to_str(value):
	'''scapy returns some record data as bytes, some as str'''
	return value.decode() if isinstance(value, bytes) else str(value)

if __name__=='__main__':
	# check parameters
	valid_types = ('A', 'AAAA', 'MX', 'NS')
	if len(sys.argv) == 2:
		target = sys.argv[1]
		record_type = 'A'
	elif len(sys.argv) == 3:
		target = sys.argv[1]
		record_type = sys.argv[2].upper()
	else:
		print(f"Usage: sudo {sys.argv[0]} <domain name> [{'|'.join(valid_types)}]")
		exit()

	if record_type not in valid_types:
		print(f"Unknown record type {record_type}. Choose from {'|'.join(valid_types)}")
		exit()

	#build DNS UDP packet on port 53 to Google at IP address 8.8.8.8
	ip = IP(dst='8.8.8.8')
	udp = UDP(dport=53)
	qdsec = DNSQR(qname=target, qtype=record_type)
	dns = DNS(id=100, qr=0, rd=1, qdcount=1, qd=qdsec)
	pkt = ip/udp/dns
	#pkt.show()

	#send packet using sr1 (send/receive 1 packet)
	reply = sr1(pkt, timeout=3, verbose=0)

	#evaluate reply
	if reply is None or not reply.haslayer(DNS):
		print("Invalid reply (no response from resolver)")
		exit()

	print("Packet is", reply)
	reply.show()

	dns = reply[DNS]
	print(f"{record_type} records for {target}")

	# answer section normally holds the records; NS records for a zone
	# sometimes arrive in the authority section (ns) instead
	records = list(dns.an or []) + list(dns.ns or [])
	found = False
	for rr in records:
		if rr.type == 1 or rr.type == 28:		# 1 is A, 28 is AAAA
			print(f"\t{rr.rdata}")
			found = True
		elif rr.type == 15:						# 15 is MX
			print(f"\tMail preference: {rr.preference}, Mail Exchange: {to_str(rr.exchange)}")
			found = True
		elif rr.type == 2:						# 2 is NS
			print(f"\tNameserver: {to_str(rr.rdata)}")
			found = True
		elif rr.type == 5:						# 5 is CNAME, often returned alongside A
			print(f"\tCNAME: {to_str(rr.rdata)}")

	if not found:
		print(f"\tNo {record_type} records returned (rcode {dns.rcode})")