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


class ServiceDelegate(ServiceBase) :

	_api_sessions = None

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

	def startup(self) :
		for requirement in ['keys'] :
			if requirement not in self._config or (requirement in self._config and self._config[requirement] is ''):
				self._error.append('Missing required config option ' + requirement)
				return
		keys = self._config['keys']
		self._api_sessions = [shodan.Shodan(x) for x in keys]
		self.debug('Searching Shodan for hosts...',1)


	def lookup(self,subject) :
		if self._api_sessions :
			session = next(self.get_api())
			try :
				host = session.host(subject)
				clean = self._clean_newlines(host)
				return host
			except shodan.exception.APIError as e:
				self.debug('Unable to search for ' + subject + ' on shodan: ' + str(e),1)
		pass

	def _clean_newlines(self,structure) :
		if type(structure) is dict :
			newstructure = {}
			for key in structure :
				if key :
					if structure[key] and type(structure[key]) is dict or type(structure[key]) is list :
						newkey = key.replace('\n','')
						newstructure.update({newkey:self._clean_newlines(structure[key])})
					else :
						newkey = key.replace('\n','')
						if structure[key] and type(structure[key]) is str :
							newvalue = structure[key].replace('\n','')
							newstructure.update({newkey:newvalue})
						else :
							newstructure.update({newkey:None})
			return newstructure
		if type(structure) is list :
			newstructure = []
			for item in structure :
				if type(item) is dict or type(item) is list :
					newstructure.append(self._clean_newlines(item))
				if type(item) is str :
					newstructure.append(item.replace('\n',''))
				else :
					newstructure.append(item)
			return newstructure

	def get_api(self) :
		while True :
			for x in self._api_sessions :
				yield x