from servicebase import ServiceBase
import requests
import sys
import re
import json
import urllib3
urllib3.disable_warnings()

class ServiceDelegate(ServiceBase) :

	_return_payload = {}
	_hosts = None
	_partition = 'Common'

	def get_arguments(cls) :
		return ['-cg', '--vip', 'store_true', 'Return VIP information about the subject (f5)']

	def startup(self) :
		for requirement in ['hosts','username','key'] :
			if requirement not in self._config or (requirement in self._config and self._config[requirement] == ''):
				self._error.append('Missing required config option ' + requirement)
				return
		self._hosts = self._config['hosts']
		self.debug('Logging into F5\'s and searching for hosts, this make take some time.',1)

	def lookup(self,subject) :
		self._return_payload = {}

		logincandidates = self.choose_host(subject)
		if len(logincandidates) == 0 :
			# Idk what to do here, do we give up or do we log in to all of them?
			self._error.append('Unable to find an LB to log into to check host ' + str(subject))
			# Search all
			logincandidates = {x:x for x in self._hosts}
			# Give up
			# return None

		for host in logincandidates.keys() :
			auth_header = self.get_iq_auth_header(host,self._config['username'],self._config['key'])
			if not auth_header :
				self._error.append('Unable to log in to ' + host)
				continue
			# Get all VIPs
			if 'vip' not in self._return_payload :
				vips = self.get_vip(host,auth_header);
				# If it's a VIP, return it and the host and let's be done
				if vips :
					self.debug('Got vips from ' + host,3)
					vip_match = [ x for x in vips if x['destination'].split('/' + self._partition + '/')[1].startswith(subject)]
					all_vips = [ x['destination'] for x in vips ]
					if len(vip_match) > 0 :
						selected_vip = vip_match[0]
						self._return_payload.update({'vip': selected_vip})
						pool_name = selected_vip['pool'].split('/' + self._partition + '/')[1]
						pool_info = self.get_pool(host,auth_header,pool_name)
						if pool_info :
							self._return_payload.update(pool_info)
						break
					else :
						self.debug('Unable to find VIP in VIPs from host ' + host,2)
				else :
				# Otherwise, get all nodes.
				# If it's in the nodes, we have to start all over and find where.
					# Take all vips and get all pools
						# Check if address is in the pool. If it is, return the node, pool, vip, and host.
				# Naw, I'll just give up.
					self.debug('Unable to get VIPs from ' + host,2)
		if len(self._return_payload) == 0 :
			self._return_payload = None
		return self._return_payload

	def shutdown(self) :
		return


	def get_iq_auth_header(self,host,username,password) :
		assumed_login_provider = 'tmos'
		token = None
		try :
			payload = { 'username' : username, 'password' : password, 'loginProviderName' : assumed_login_provider }
			login = requests.post('https://' + host + '/mgmt/shared/authn/login', json=payload,verify=False,timeout=2)
			if 'token' in json.loads(login.text) :
				token = json.loads(login.text)['token']['token']
		except Exception as exception:
			self.debug('Exception getting auth header for host ' + host + ': \n' + str(exception),2)
			return None
		if token :
			return { 'X-F5-Auth-Token' : token }
		else :
			return None

	def get_vip(self,host,header) :
		"""Given a host and an authorization header, gets all the vips for that host"""
		# Please don't give me empty auth headers, I don't want to check it everywhere
		self.debug('Trying VIP info for ' + host,2)
		try :
			vips = requests.get('https://' + host + '/mgmt/tm/ltm/virtual',headers=header,verify=False,timeout=10)
			if vips.status_code == 200 :
				return json.loads(vips.text)['items']
			self.debug('Response not OK for ' + vips.history[0].url,3)
			return None
		except Exception as exception :
			self.debug('Exception getting VIPs for host ' + host + ': \n' + str(exception),2)
			return None

	def get_pool(self,host,header,pool) :
		"""Given a host, auth header, and pool, return the pool and it's members from the host"""
		try :
			pools = requests.get('https://' + host + '/mgmt/tm/ltm/pool/~' + self._partition + '~' + pool,headers=header,verify=False,timeout=5)
			if pools.status_code == 200 :
				poolobj = json.loads(pools.text)
				members = requests.get('https://' + host + '/mgmt/tm/ltm/pool/~' + self._partition + '~' + pool + '/members',headers=header,verify=False,timeout=10)
				if members.status_code == 200 :
					membersobj = json.loads(members.text)['items']
					return {'pool':poolobj,'members':membersobj}
				return {'pool':poolobj}
			elf.debug('Response not OK for ' + pool.history[0].url,3)
			return None
		except Exception as exception:
			self.debug('Exception getting pool for host ' + host + ': \n' + str(exception),2)
			return None

	def choose_host(self,subject) :
		"""Given a subject, return the host that we suspect is responsible for the VIP resides in."""
		hosts = {}
		# Guess based on hostname
		dcmatch = re.search('^[a-zA-Z]*-([a-zA-Z0-9]*)(-.*)*$',subject)
		if dcmatch :
			potentialdc = dcmatch.group(1)
		else :
			potentialdc = None
		if potentialdc :
			for host in self._hosts.keys() :
				if potentialdc.lower() in host.lower() :
					hosts.update({host:self._hosts[host]})
			if len(hosts) > 0 :
				return hosts
		self.debug('Unable to determine DC from subject name: ' + subject,3)
		# No hostname match, maybe we have info from bluecat?
		if 'bluecat' in self.hints and subject in self.hints['bluecat'] and self.hints['bluecat'][subject]['object'] :
			if 'locationCode' in self.hints['bluecat'][subject]['object'] :
				potentialdc = self.hints['bluecat'][subject]['object']['locationCode'].split(' ')[-1]
			elif 'parent' in self.hints['bluecat'][subject] and 'name' in self.hints['bluecat'][subject]['parent'] :
				potentialdc = self.hints['bluecat'][subject]['parent']['name'].split(' ')[0]
			for host in self._hosts.keys() :
				if potentialdc.lower() in host.lower() :
					hosts.update({host:self._hosts[host]})
			self.debug('Got additional host info from the BlueCat plugin!',3)
			return hosts
		return hosts