#!/usr/bin/python
"""This performs simultaneous dns lookups for pairs of   dns_object query_type    provided on stdin."""

from __future__ import print_function

import sys
import pprint
import argparse
#try:
#	import adns
#except ImportError:
#	#cd /opt/local/Library/Frameworks/Python.framework/Versions/2.7/include/python2.7/ ; sudo ln -sf /opt/local/include/adns.h adns.h ; cd -
#	print("Missing adns module; perhaps 'sudo rpm -ivh adns-python-1.2.2-1.x86_64.rpm', 'sudo -H pip install https://storage.googleapis.com/google-code-archive-downloads/v2/code.google.com/adns-python/adns-python-1.2.1.tar.gz' or 'sudo port install adns' ?  Exiting.")
#	raise
from async_resolver_class import AsyncResolver


def debug(debug_message):
	"""Print a debugging message."""
	sys.stderr.write('___ ' + debug_message + '\n')
	sys.stderr.flush()


def fail(debug_message):
	"""Print a debugging message and quit."""
	sys.stderr.write('xxx ' + debug_message + '\n')
	sys.stderr.flush()
	sys.exit(1)


def show_answers(answer_dict, do_prettyprint, do_typesummary, do_csv):
	if do_prettyprint:
		pp = pprint.PrettyPrinter(indent=4)
		pp.pprint(answer_dict)

	if do_typesummary:
		for one_dns_obj in sorted(answer_dict):
			live_answers = []
			for one_ans_type in sorted(answer_dict[one_dns_obj]):
				if answer_dict[one_dns_obj][one_ans_type] != None:
					live_answers.append(one_ans_type)
			print(one_dns_obj + ': ' + '/'.join(live_answers))

	if do_csv:
		for one_dns_obj in sorted(answer_dict):
			for one_ans_type in sorted(answer_dict[one_dns_obj]):
				if answer_dict[one_dns_obj][one_ans_type] != None:
					for one_answer in answer_dict[one_dns_obj][one_ans_type]:
						clean_answer = str(one_answer).replace(',', '_')
						clean_dns_obj = str(one_dns_obj).replace(',', '_')
						if clean_dns_obj[-1] != '.':
							clean_dns_obj = clean_dns_obj + '.'
						if one_ans_type == 'A':
							print('DN,' + clean_answer + ',A,' + clean_dns_obj + ',')
						elif one_ans_type == 'CNAME':
							#This record only has the CNAME names, not the target IP needed by passer.  Skip it and depend on A records instead.
							pass
							#print('DN,0.0.0.0,CNAME,' + clean_dns_obj + ',' + clean_answer)
						elif one_ans_type == 'MX':
							#This record only has the mailserver names.  If we have an A record for the mailserver name, now we can create full MX record lines.
							if clean_answer in answer_dict and 'A' in answer_dict[clean_answer] and answer_dict[clean_answer]['A']:
								for one_mx_ip in answer_dict[clean_answer]['A']:
									print('DN,' + one_mx_ip + ',MX,' + clean_dns_obj + ',' + clean_answer)
						elif one_ans_type == 'MX_IP':
							print('DN,' + clean_answer + ',MX,' + clean_dns_obj + ',')
						elif one_ans_type == 'NS':
							#This record only has the nameserver names.  If we have an A record for the nameserver name, now we can create full NS record lines.
							if clean_answer in answer_dict and 'A' in answer_dict[clean_answer] and answer_dict[clean_answer]['A']:
								for one_ns_ip in answer_dict[clean_answer]['A']:
									print('DN,' + one_ns_ip + ',NS,' + clean_dns_obj + ',' + clean_answer)
						elif one_ans_type == 'NS_IP':
							print('DN,' + clean_answer + ',NS,' + clean_dns_obj + ',')
						elif one_ans_type == 'SOA':
							print('DN,0.0.0.0,SOA,' + clean_dns_obj + ',' + clean_answer)
						elif one_ans_type == 'TXT':
							print('DN,0.0.0.0,TXT,' + clean_dns_obj + ',' + clean_answer)
						else:
							fail('Unrecognized one_ans_type in show_answers: ' + str(one_ans_type))

#Main
slice_size = 500		#This should be a multiple of 5 to avoid having answers for a given domain should up at the end of one slice and the beginning of the next.
mega_resolve_ver = '1.6'

parser = argparse.ArgumentParser(description='Mega_resolve v' + mega_resolve_ver + '. Lookup large numbers of dns objects, provided as   dns_obj query_type   pairs on stdin.')
parser.add_argument('-p', '--pprint', '--prettyprint', help='Show output in a pretty printed tree', dest='prettyprint', action='store_true')
parser.add_argument('-t', '--typesummary', help='Summarize which dns types are available for each object', dest='typesummary', action='store_true')
parser.add_argument('-c', '--csv', help='Show all answers in csv format', dest='csv', action='store_true')
parser.add_argument('-s', '--simul', help='Maximum simultaneous DNS requests', required=False)
parser.set_defaults(prettyprint=False,typesummary=False,csv=False)
args = vars(parser.parse_args())

prettyprint = args['prettyprint']
typesummary = args['typesummary']
csv = args['csv']
if not prettyprint and not typesummary and not csv:
	typesummary = True

if args['simul']:
	max_intensity = args['simul']
else:
	max_intensity = 100

dom_requests = []

for one_line in sys.stdin:
	req_atoms = one_line.replace('\n', '').split()
	if len(req_atoms) == 0:
		pass
	elif len(req_atoms) == 1:
		dom_requests.append((req_atoms[0], 'A'))
		dom_requests.append((req_atoms[0], 'MX'))
		dom_requests.append((req_atoms[0], 'NS'))
		dom_requests.append((req_atoms[0], 'SOA'))
		dom_requests.append((req_atoms[0], 'TXT'))
	elif len(req_atoms) == 2:
		if req_atoms[1] in ('A', 'MX', 'NS', 'SOA', 'TXT'):
			dom_requests.append((req_atoms[0], req_atoms[1]))
		elif req_atoms[1] == 'ANY':
			dom_requests.append((req_atoms[0], 'A'))
			dom_requests.append((req_atoms[0], 'MX'))
			dom_requests.append((req_atoms[0], 'NS'))
			dom_requests.append((req_atoms[0], 'SOA'))
			dom_requests.append((req_atoms[0], 'TXT'))
		else:
			fail('Invalid dns query type for line: ' + one_line.replace('\n', ''))
	elif len(req_atoms) > 2:
		fail('Invalid format for line: ' + one_line.replace('\n', ''))

	if len(dom_requests) >= slice_size:
		dom_request_slice = dom_requests[0:slice_size]
		dom_requests = dom_requests[slice_size:]
		ar = AsyncResolver(dom_request_slice, intensity=max_intensity)
		all_answers = ar.resolve()

		show_answers(all_answers, prettyprint, typesummary, csv)

if len(dom_requests) > 0:
	ar = AsyncResolver(dom_requests, intensity=max_intensity)
	all_answers = ar.resolve()

	show_answers(all_answers, prettyprint, typesummary, csv)