#!/usr/bin/env python3

# - example usage : python3 parse-rr-bin.py rr2.bin 22 -c 12

import argparse
import re
import struct
from dataclasses import dataclass

DEFAULT_COMPRESSION_BASE=12

MDNS_TYPE = {
	1:	'A',	# RFC 1035 3.2.2
	12:	'PTR',	# RFC 1035 3.2.2
	16:	'TXT',	# RFC 1035 3.2.2
	28:	'AAAA',	# RFC 3596
	33:	'SRV',	# RFC 2782
	10:	'NULL', # RFC 1035 3.2.2
}

@dataclass
class ResourceRecord:
	""" mDNS Resource Record conforming to RFC 1035 """
	container: bytearray
	offset: int

	# fields which are populated after evaluate
	_my_bytes: bytearray = None
	_name: str = None
	_name_bytes_used: int = 0
	_type: int = 0
	_class: int = 0
	_ttl: int = 0
	_rd_length: int = 0
	_rd_data: bytearray = None

	def __str__(self):
		_my_str = ''
		_my_str += f'name: {self._name}\n'
		_my_str += f'type: {self._type}\n'
		_my_str += f'class: {self._class}\n'
		_my_str += f'ttl: {self._ttl}\n'
		_my_str += f'rd_length: {self._rd_length}\n'
		return(_my_str)
	
	def get_bytes_used(self):
		""" Compute the number of bytes that this RR used """
		bytes_used = 0
		bytes_used += self._name_bytes_used
		bytes_used += (2 + 2 + 4 + 2) # HHIH
		bytes_used += self._rd_length
		
		return bytes_used

	def parse_name(self, buffer):
		""" Traverse a buffer, and parse names according to RFC 1035 3.2.1 and 4.1.4 (compression) """
		buffer_idx = 0
		ret_name_str = ''
		bytes_used = 0

		len0 = buffer[buffer_idx]
		# keep reading bytes until we get to the zero length trailer
		while (int(len0) != 0):
			if ((int(len0) & int('0xc0', 0)) == 0):
				# a normal non-compressed string
				name_bytes = buffer[buffer_idx + 1:buffer_idx + 1 + len0]
				ret_name_str += name_bytes.decode(encoding="ascii")
				ret_name_str += '.'

				buffer_idx = buffer_idx + 1 + len0
				bytes_used += 1 + len0

				len0 = buffer[buffer_idx]
			else:
				# a compressed string
				len1 = buffer[buffer_idx + 1]
				next_offset = ((len0 & 0x3F) << 8) | len1

				# recursively call to get our answer
				(compressed_name_str, compressed_bytes_used) = self.parse_name(self.container[next_offset:])
				ret_name_str += compressed_name_str

				# but, for the sake of offsets, we've only consumed 2 bytes (len0, len1)
				bytes_used += 2

				break

		# need to account for the zero length trailer
		if (int(len0) == 0):
			bytes_used += 1

		return (ret_name_str, bytes_used)

	def evaluate(self):
		""" Populate fields of this RR from the containing buffer """
		self._my_bytes = self.container[self.offset:]

		(self._name, self._name_bytes_used) = self.parse_name(self._my_bytes)
		(self._type, self._class, self._ttl, self._rd_length) = struct.unpack_from('!HHIH', self._my_bytes, self._name_bytes_used)
		rd_offset = self.get_bytes_used() - self._rd_length
		self._rd_data = self._my_bytes[rd_offset:rd_offset + self._rd_length]

def decode_rr_rdata_srv(rr):
	# RFC 2782
	(priority, weight, port) = struct.unpack_from('!HHH', rr._rd_data)
	(target, target_bytes_used) = rr.parse_name(rr._rd_data[6:])
	return('\t SRV prio {:5d}, weight {:5d}, port {:5d}, target {}'.format(priority, weight, port, target))

def decode_rr_rdata_null(rr):
	# RFC 1035 3.2.10 states that the RDATA format that follows is experimental
	# for Apple, the NULL RDATA shall have different formats depending on IPv4 or IPV6:
	# RDATA[0]    : length of keepalive text which follows
	# RDATA[1...] : keepalive text = "t=%d i=%d c=%d h=%#a d=%#a l=%u r=%u m=%.6a s=%u a=%u w=%u"
	# RDATA[1...] : keepalive text = "t=%d i=%d c=%d H=%#a D=%#a l=%u r=%u m=%.6a s=%u a=%u w=%u"
	# IPv4 and IPv6 are different based on 'h=,d=' vs. 'H=,D='

	tcpka_text_len = rr._rd_data[0]
	tcpka_text = rr._rd_data[1:].decode('ascii')

	tcpka_ipv4_re = re.compile('t=([0-9]+) i=([0-9]+) c=([0-9]+) h=(.*) d=(.*) l=([0-9]+) r=([0-9]+) m=(.*) s=([0-9]+) a=([0-9]+) w=([0-9]+)')
	tcpka_ipv6_re = re.compile('t=([0-9]+) i=([0-9]+) c=([0-9]+) H=(.*) D=(.*) l=([0-9]+) r=([0-9]+) m=(.*) s=([0-9]+) a=([0-9]+) w=([0-9]+)')

	tcpka_ipv4 = tcpka_ipv4_re.match(tcpka_text)
	tcpka_ipv6 = tcpka_ipv6_re.match(tcpka_text)

	if (tcpka_ipv4):
		use_match = tcpka_ipv4
	elif (tcpka_ipv6):
		use_match = tcpka_ipv6
	else:
		return '\tUnexpected RDATA format'

	timeout = use_match.group(1)
	retry_interval = use_match.group(2)
	retry_count = use_match.group(3)
	local_addr = use_match.group(4)
	remote_addr = use_match.group(5)
	local_port = use_match.group(6)
	remote_port = use_match.group(7)
	remote_mac = use_match.group(8)
	tcp_seq = use_match.group(9)
	tcp_ack = use_match.group(10)
	tcp_win = use_match.group(11)

	return('\t' + tcpka_text)

def main():
	parser = argparse.ArgumentParser(description='Parse binary file for mDNS Resource Records')
	parser.add_argument('filename', help='input binary file', type=argparse.FileType('rb'))
	parser.add_argument('num_records', help='number of expected Resource Records', type=int)
	parser.add_argument('-c', '--compression', default=DEFAULT_COMPRESSION_BASE, help=f'offset of the first RR, default {DEFAULT_COMPRESSION_BASE}')
	parser.add_argument('-v', '--verbose', action='store_true')
	args = parser.parse_args()

	raw_buffer = args.filename.read()

	cur_offset = int(args.compression)

	all_rr=[]
	for i in range(int(args.num_records)):
		rr = ResourceRecord(container=raw_buffer, offset=cur_offset)
		rr.evaluate()
		all_rr.append(rr)
		cur_offset += rr.get_bytes_used()

	print('{:4s} {:64s} {:4s} {:6s} {:8s} {:6s}'.format(
		'idx',
		'name',
		'type',
		'class',
		'ttl',
		'rd_len'))
	print(100*'-')

	i = 0
	for rr in all_rr:
		friendly_type = MDNS_TYPE[int(rr._type)]
		if (friendly_type == 'SRV'):
			rr_details = decode_rr_rdata_srv(rr)
		elif (friendly_type == 'NULL'):
			rr_details = decode_rr_rdata_null(rr)
		else:
			rr_details = ''
		print('{:4d} {:64s} {:4s} 0x{:04x} {:8d} 0x{:04x} {}'.format(
			i,
			rr._name,
			friendly_type,
			rr._class,
			rr._ttl,
			rr._rd_length,
			rr_details))

		i += 1
	return

if __name__ == "__main__" :
	main()
