# 
# This is a ruby DNS packet encoder/decoder 
# the sample script at the end shows a simple name resolution
#
# author: Yoann Guillot (2008)
# license: wtfpl
#

class ResourceRecord
	TYPE = %w[0 A NS MD MF CNAME SOA MB MG MR NULL WKS PTR HINFO MINFO MX TXT] +
	       %w[RP AFSDB X25 ISDN RT NSAP NSAPPTR SIG KEY PX GPOS AAAA LOC NXT EID] +
	       %w[NIMLOC SRV ATMA NAPTR KX CERT A6 DNAME SINK OPT UINFO UID GID]
	TYPE[252] = 'AXFR'
	TYPE[253] = 'MAILB'
	TYPE[254] = 'MAILA'
	TYPE[255] = '*'

	DNSCLASS = %w[0 IN CS CH HS]
	DNSCLASS[255] = '*'

	def self.decode_query(p) new.decode_query(p) end
	def self.decode(p) new.decode(p) end
	def self.new_q(name, type, cls='IN')
		q = new
		q.name = name
		q.type = type
		q.dnsclass = cls
		q
	end
	def self.new_r(name, type, rdata, ttl=3600, cls='IN')
		r = new_q(name, type, cls)
		r.ttl = ttl
		r.rdata = rdata
		r
	end

	attr_accessor :name, :type, :dnsclass, :ttl, :rdlength, :rdata_raw, :rdata
	def decode_query(dnspacket)
		@name = dnspacket.decode_domain
		@type = dnspacket.readshort
		@type = TYPE[@type] || @type
		@dnsclass = dnspacket.readshort
		@dnsclass = DNSCLASS[@dnsclass] || @dnsclass
		self
	end

	def decode(dnspacket)
		decode_query(dnspacket)
		@ttl = dnspacket.readlong
		@rdata_raw = dnspacket.read(dnspacket.readshort)
		interpret(dnspacket)
		self
	end

	def interpret(dnspacket)
		str = @rdata_raw.dup
		@rdata = case @type
		when 'CNAME', 'PTR', 'NS': dnspacket.decode_domain(str)
		when 'HINFO': { :cpu => dnspacketstr.readstr(str), :os => dnspacket.readstr(str) }
		when 'MX': { :preference => dnspacket.readshort(str), :exchange => dnspacket.decode_domain(str) }
		when 'SOA': { :mname => dnspacket.decode_domain(str), :rname => dnspacket.decode_domain(str),
			:serial => dnspacket.readlong(str), :refresh => dnspacket.readlong(str), :retry => dnspacket.readlong(str), 
			:expire => dnspacket.readlong(str), :minimum => dnspacket.readlong(str) }
		when 'A': dnspacket.read(4, str).unpack('C*').join('.')
		when 'AAAA': dnspacket.read(16, str).unpack('n*').map { |i| i.to_s 16 }.join(':')
		end
	end

	def encode_query(dnspacket)
		dnspacket.encode_domain(@name)
		dnspacket.writeshort(TYPE.index(@type) || @type)
		dnspacket.writeshort(DNSCLASS.index(@dnsclass || 'IN') || @dnsclass)
	end
	def encode(dnspacket)
		encode_query(dnspacket)
		dnspacket.writelong(@ttl)
		encode_rdata(dnspacket)
		dnspacket.writeshort(@rdata_raw.length)
		dnspacket.write(@rdata_raw)
	end

	def encode_rdata(dnspacket)
		@rdata_raw = ''
		return if @rdata.empty?
		case @type
		when 'CNAME', 'NS', 'PTR': dnspacket.encode_domain(@rdata, @rdata_raw)
		when 'HINFO': dnspacket.writestr(@rdata[:cpu], @rdata_raw) ; dnspacket.writestr(@rdata[:os], @rdata_raw)
		when 'MX': dnspacket.writeshort(@rdata[:preference], @rdata_raw) ; dnspacket.encode_domain(@rdata[:exchange], @rdata_raw)
		when 'SOA': dnspacket.encode_domain(@rdata[:mname], @rdata_raw) ; dnspacket.encode_domain(@rdata[:rname], @rdata_raw)
			[:serial, :refresh, :retry, :expire, :minimum].each { |s| dnspacket.writelong(@rdata[s], @rdata_raw) }
		when 'A': dnspacket.write(@rdata.split('.').map { |i| i.to_i }.pack('C*'), @rdata_raw)
		when 'AAAA':
			bits = @rdata.split(':')
			if i = bits.index('')
				bits[i, 1] = %[0] * (9-bits.length)
			end
			dnspacket.write(bits.map { |i| i.to_i(16) }.pack('n*'), @rdata_raw)
		end
	end

	def to_s
		a = ["#{@name.inspect} #@type #@dnsclass#{" TTL #@ttl" if @ttl}"]
		a << (@rdata || @rdata_raw).inspect if @rdata or @rdata_raw
		a.join(" ")
	end
end

class DNSPacket
	class Pointer
		attr_accessor :ptr
		def initialize(ptr) @ptr = ptr end
	end

	OPCODE = %w[QUERY IQUERY STATUS]
	RCODE = %w[NOERROR FMTERR SERVFAIL NXDOMAIN NOTIMPLEMENTED REFUSED]

	def read(len, str=@curstr)
		ret = str[0, len]
		str[0, len] = ''
		ret
	end
	def readstr(str=@curstr)
		return '' if not len = read(1, str)[0]
		if (len >> 6) == 0b11
			return '' if str.empty?
			len = len & 0b11_1111
			len = (len << 8) | read(1, str)[0].to_i
			Pointer.new(len)
		else
			read(len, str)
		end
	end
	def readshort(str=@curstr)
		read(2, str).unpack('n').first
	end
	def readlong(str=@curstr)
		read(4, str).unpack('N').first
	end

	def write(str, out=@curstr)
		if out.equal? @curstr
			@domain_cache.each { |k, v|
				if v.kind_of? Array and v[0] == str
					@domain_cache[k] = out.length + v[1]
				end
			}
		end
		out << str
	end
	def writestr(str, out=@curstr)
		out << str.length << str
	end
	def writeshort(i, out=@curstr)
		out << [i].pack('n')
	end
	def writelong(i, out=@curstr)
		out << [i].pack('N')
	end

	def decode_domain(str=@curstr, recurs=false)
		ret = []
		loop do
			case s = readstr(str)
			when Pointer
				if recurs and recurs <= 0
					ret << '<recursion>'
				else
					ret.concat decode_domain(@fullpacket[s.ptr..-1], (recurs ? recurs-1 : 8))
				end
				break
			when '': break
			when String: ret << s
			end
		end
		recurs ? ret : ret.join('.')
	end

	def encode_domain(dom, out=@curstr)
		dom = dom.split('.')
		until dom.empty?
			if @domain_cache[dom].kind_of? Integer
				writeshort((0b11 << 14) | @domain_cache[dom])
				return
			end
			@domain_cache[dom] ||= (out.equal? @curstr ? out.length : [out, out.length])
			writestr(dom.shift, out)
		end
		writestr('', out)
	end

	def self.decode(p) new.decode(p) end

	attr_accessor :id, :qr, :opcode, :aa, :tc, :rd, :ra, :z, :rcode
	attr_accessor :question, :answer, :authority, :additional
	def decode(fullpacket)
		@fullpacket = fullpacket
		@curstr = fullpacket.dup

		@id = readshort
		flags = readshort
		@qr = (flags[15] == 0 ? :query : :response)
		@opcode = (flags >> 11) & 0b1111
		@opcode = OPCODE[@opcode] || @opcode
		@aa = (flags[10] == 1 ? :authoritative : nil)
		@tc = (flags[9] == 1  ? :truncated : nil)
		@rd = (flags[8] == 1  ? :recursion_desired : nil)
		@ra = (flags[7] == 1  ? :recursion_available : nil)
		@z = (flags >> 4) & 0b111
		@z = nil if @z == 0
		@rcode = flags & 0b1111
		@rcode = RCODE[@rcode] || @rcode

		qdcount = readshort
		ancount = readshort
		nscount = readshort
		arcount = readshort
		@question   = (0...qdcount).map { ResourceRecord.new.decode_query(self) }
		@answer     = (0...ancount).map { ResourceRecord.new.decode(self) }
		@authority  = (0...nscount).map { ResourceRecord.new.decode(self) }
		@additional = (0...arcount).map { ResourceRecord.new.decode(self) }
		self
	end

	def encode
		@curstr = ''
		@domain_cache = {}
		flags = 0
		flags |= 1 << 15 if @qr == :response
		op = OPCODE.index(@opcode || 'QUERY') || @opcode
		flags |= op.to_i << 11
		flags |= 1 << 10 if @aa
		flags |= 1 << 9  if @tc
		flags |= 1 << 8  if @rd
		flags |= 1 << 7  if @ra
		flags |= @z.to_i << 4
		rc = RCODE.index(@rcode || 'NOERROR') || @rcode
		flags |= rc.to_i

		@id ||= rand(0x10000)
		[@id, flags, @question.to_a.length, @answer.to_a.length, @authority.to_a.length, @additional.to_a.length].each { |s| writeshort(s) }
		@question.to_a.each { |q| q.encode_query(self) }
		(@answer.to_a+@authority.to_a+@additional.to_a).each { |q| q.encode(self) }
		@curstr
	end

	def add_question(q)   (@question   ||= []) << q ; self end
	def add_answer(q)     (@answer     ||= []) << q ; self end
	def add_authority(q)  (@authority  ||= []) << q ; self end
	def add_additional(q) (@additional ||= []) << q ; self end

	def to_s
		a = [[('DNS id=%04X' % @id), @qr, @opcode, @aa, @tc, @rd, @ra, @z, @rcode].compact.join(' ')]
		a << 'question:'   << @question.join("\n")   if @question.to_a   != []
		a << 'answer:'     << @answer.join("\n")     if @answer.to_a     != []
		a << 'authority:'  << @authority.join("\n")  if @authority.to_a  != []
		a << 'additional:' << @additional.join("\n") if @additional.to_a != []
		a.join("\n")
	end
end

class DNSResolver
	def initialize(sock)
		@sock = sock || UDPSocket.open
		@send_cb = proc { |ip, rq|
			@sock.send(rq.encode, 0, ip, 53)
			DNSPacket.decode @sock.recv(512) if IO.select([@sock], nil, nil, 3)
		}

		@cache = {
			['NS', 'com'] => 'a.gtld-servers.net',
			['A', 'a.gtld-servers.net'] => '192.5.6.30',
		}
	end

	def response(request)
		ans = DNSPacket.new
		ans.id = request.id
		ans.opcode = request.opcode
		ans.question = request.question
		ans.qr = :response
		ans.rd = true
		ans.ra = true
		ans.question.each { |q| resolve(q, ans) }
		ans.rcode = 'NXDOMAIN' if ans.response.to_a.empty?
		ans
	end

	def resolve(q, ans)
		case q.type
		when 'A', 'CNAME'
			recurse(q.name) if not @cache[['A', q.name]] and not @cache[['CNAME', q.name]]
			if a = @cache[['A', q.name]]
				ans.add_answer ResourceRecord.new_r(q.name, 'A', a, 10)
			elsif cn = @cache[['CNAME', q.name]]
				done = []
				until done.include? cn
					recurse(cn) if not @cache[['A', cn]]
					break if not @cache[['CNAME', cn]]
					done << cn
					cn = @cache[['CNAME', cn]]
				end
				ans.add_answer ResourceRecord.new_r(q.name, 'CNAME', cn, 10)
				if a = @cache[['A', cn]]
					ans.add_additional ResourceRecord.new_r(cn, 'A', a, 10)
				end
			end
		when 'NS'
			recurse('ns.' + q.name) if not @cache[['NS', q.name]]
			if ns = @cache[['NS', q.name]]
				ans.add_answer ResourceRecord.new_r(q.name, 'NS', ns, 10)
				if a = @cache[['A', ns]]
					ans.add_additional ResourceRecord.new_r(ns, 'A', a, 10)
				end
			end
		when 'PTR'
			#recurse(q.name) if not @cache[[q.type, q.name]]
		end
	end

	def recurse(name)
		name_a = name.split('.')
		return if not nsip = (1...name_a.length).map { |l| @cache[['NS', name_a[l..-1].join('.')]] }.compact.first
		rq = DNSPacket.new.add_question ResourceRecord.new_q(name, 'A')
		return if not a = @send_cb[nsip, rq]
		(a.answer + a.authority + a.additional).each { |ans|
			@cache[[ans.type, ans.name]] = ans.rdata
		}
		recurse(name) if a.answer.empty? and a.rcode == 'NOERROR'
	end
end

class DNSServer
	def initialize(domain='me.com', ip='127.0.0.1', port=53)
		@sock = UDPSocket.open
		@sock.bind(ip, port)
		@myip = ip
		@domain = domain
	end

	def mainloop
		loop do
			next if not IO.select([@sock], nil, nil, 1)
			begin
				packet = @sock.recvfrom(512)
				rq = DNSPacket.decode(packet[0])
				puts "\n<-- #{packet[1][3]} #{packet[1][1]} #{rq.id} #{rq.question.first.type} #{rq.question.first.name}"
File.open('dns.log', 'a') { |fd| fd.puts "#{packet[1][3]} #{packet[1][1]} #{rq.id} #{rq.question.first.type} #{rq.question.first.name}" }
				if ans = respond(rq)
					puts '-->', ans
					@sock.send(ans.encode, 0, packet[1][3], packet[1][1])
				end
			rescue Exception
				puts $!, $!.backtrace
			end
		end
	end

	def respond(request)
		ans = DNSPacket.new
		ans.id = request.id
		ans.rd = request.rd
		ans.opcode = request.opcode
		ans.question = request.question
		ans.qr = :response
		ans.aa = :authoritative
		ans.question.each { |q|
			case q.type
			when 'A'
				respond_a(ans, q)
				ans.rcode = 'NXDOMAIN' if ans.answer.to_a.empty?
			when 'NS'
				ans.add_answer ResourceRecord.new_r(q.name, q.type, 'ns.'+@domain, 30)
				ans.add_additional ResourceRecord.new_r('ns.'+@domain, 'A', @myip, 30)
			end
		}
		ans
	end

	def respond_a(ans, q)
		case q.name
		when 'poison.'+@domain
			ans.add_answer ResourceRecord.new_r(q.name, q.type, '1.2.3.4', 10)
			ans.add_authority ResourceRecord.new_r(@domain, 'NS', 'evilns.'+@domain, 30)
			ans.add_additional ResourceRecord.new_r('evilns.'+@domain, 'A', '5.6.7.8', 30)
		else	
			ans.add_answer ResourceRecord.new_r(q.name, q.type, @myip, 10)
		end
	end
end

if $0 == __FILE__

require 'socket'

Socket.do_not_reverse_lookup = true

# launch a self-poisonning dns server 
#DNSServer.new.mainloop

name = ARGV.shift || 'www.l.google.com'
serv = ARGV.shift || File.read('/etc/resolv.conf').grep(/nameserver/i).first.split.last
servip, servport = serv.split ':'
servport ||= 53
type = ARGV.shift || 'A'

# build request
packet = DNSPacket.new
packet.opcode = 'QUERY'
packet.rd = true
packet.id = rand(0x10000)
packet.add_question ResourceRecord.new_q(name, type)
raw = packet.encode

# show request
puts '-'*30 + ' query  ' + '-'*30, raw.inspect, '', DNSPacket.decode(raw)

# send request
s = UDPSocket.open
s.connect servip, servport
s.send raw, 0

# read response
resp = s.recv(512)

# show response
puts '-'*30 + ' answer ' + '-'*30, resp.inspect, '', DNSPacket.decode(resp)

end