from servicebase import ServiceBase
import socket
import ipaddress
import dns.resolver


class ServiceDelegate(ServiceBase) :

	def get_arguments(cls) :
		"""Returns an array of information used to construct an argumentparser argument."""
		return [ '-d', '--dns', 'store_true', 'Return DNS resolution about the subject (dns)' ]

	def startup(self) :
		self.resolver = dns.resolver.Resolver()
		if 'resolver' in self._config :
			self.resolver.nameservers.append(self._config['resolver'])

	def lookup(self,subject) :

		# Note : some hostname lookups always use system resolver.
		try :
			subjecttuple = socket.gethostbyname_ex(subject)
		# Can't make an IP of the subject
		except ValueError :
			try :
				# Lookup with system
				subjecttuple = socket.gethostbyname_ex(subject)
			except socket.gaierror :
				# Not a valid host or IP
				self._error.append('Unable to resolve ' + subject)
				subjecttuple = [subject,[],['']]
		# Can't make IPv4 out of it
		except socket.gaierror :
			self._error.append('Unable to resolve ' + subject)
			subjecttuple = [subject,[],['']]
		# Get additional records
		return_dict = {}
		for record_type in ['A','AAAA','CNAME','MX','NS','TXT'] :
			recs = sorted(self.get_records(subject,record_type))
			if recs and len(recs) > 0 :
				return_dict.update({record_type.lower():recs})
		# Combine Addresses
		if len(subjecttuple[2]) > 0 and subjecttuple[2][0] != '':
			return_dict['addresses'] = list(subjecttuple[2])
		if 'a' in return_dict :
			return_dict['addresses'] += [x for x in return_dict['a'] if x != subjecttuple[2]]
		if 'aaaa' in return_dict :
			return_dict['addresses'] += [x for x in return_dict['aaaa'] if x != subjecttuple[2]]
		if 'addresses' in return_dict :
			return_dict['addresses'] = list(set(return_dict['addresses']))
		# Include name and aliases
		if subjecttuple[0] and subjecttuple[0] != '' and len(return_dict) > 0:
			return_dict['name'] = subjecttuple[0]
		if len(subjecttuple[1]) > 0 :
			return_dict['aliases'] = list(subjecttuple[1])
		if len(return_dict) > 0 :
			return return_dict
		else :
			return None


	def get_records(self,subject,recordtype) :
		recs = []
		try :
			results = self.resolver.query(subject,recordtype)
			for result in results :
				recs.append(result.to_text().strip('\"'))
			return recs
		except Exception as exception :
			return  recs