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] is ''):
				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 perform_lookup(self,subject) :
		self.return_payload = {}

		logincandidates = self.choose_host(subject)
		if logincandidates == {} :
			# 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))
			return self.return_payload

		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 self.return_payload == {} :
			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 hosts != {} :
				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.data and subject in self.data['bluecat'] and self.data['bluecat'][subject]['object'] :
			if 'locationCode' in self.data['bluecat'][subject]['object'] :
				potentialdc = self.data['bluecat'][subject]['object']['locationCode'].split(' ')[-1]
			elif 'parent' in self.data['bluecat'][subject] and 'name' in self.data['bluecat'][subject]['parent'] :
				potentialdc = self.data['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