from servicebase import ServiceBase
import requests
from xml.etree.ElementTree import Element, SubElement, Comment, tostring
import xml.etree.ElementTree as ElementTree


class ServiceDelegate(ServiceBase) :

	host = None
	api_session = requests.session()
	session_key = None
	query_page_size = 200
	clients_fields = ['sta_mac_address', 'client_ht_phy_type', 'openflow_state', 'client_ip_address', 'client_user_name', 'client_dev_type', 'client_ap_location', 'client_conn_port', 'client_conn_type', 'client_timestamp', 'client_role_name', 'client_active_uac', 'client_standby_uac', 'ap_cluster_name', 'client_health', 'total_moves', 'successful_moves', 'steer_capability', 'ssid', 'ap_name', 'channel', 'channel_str', 'channel_busy', 'tx_time', 'rx_time', 'channel_free', 'channel_interference', 'current_channel_utilization', 'radio_band', 'bssid', 'speed', 'max_negotiated_rate', 'noise_floor', 'radio_ht_phy_type', 'snr', 'total_data_frames', 'total_data_bytes', 'avg_data_rate', 'tx_avg_data_rate', 'rx_avg_data_rate', 'tx_frames_transmitted', 'tx_frames_dropped', 'tx_bytes_transmitted', 'tx_bytes_dropped', 'tx_time_transmitted', 'tx_time_dropped', 'tx_data_transmitted', 'tx_data_dropped', 'tx_data_retried', 'tx_data_transmitted_retried', 'tx_data_bytes_transmitted', 'tx_abs_data_bytes', 'tx_data_bytes_dropped', 'tx_time_data_transmitted', 'tx_time_data_dropped', 'tx_mgmt', 'rx_frames', 'rx_bytes', 'rx_data', 'rx_data_bytes', 'rx_abs_data_bytes', 'rx_data_retried', 'tx_data_frame_rate_dist', 'rx_data_frame_rate_dist', 'tx_data_bytes_rate_dist', 'rx_data_bytes_rate_dist', 'connection_type_classification', 'total_data_throughput', 'tx_data_throughput', 'rx_data_throughput', 'client_auth_type', 'client_auth_subtype', 'client_encrypt_type', 'client_fwd_mode']
	aps_fields = ['mon_ap', 'mon_bssid', 'mon_radio_phy_type', 'mon_ssid', 'mon_radio_band', 'mon_ap_current_channel', 'mon_ht_sec_channel', 'mon_sta_count', 'mon_ap_classification', 'mon_ap_match_conf_level', 'mon_ap_encr', 'mon_ap_encr_auth', 'mon_ap_encr_cipher', 'mon_ap_is_dos', 'mon_ap_type', 'mon_ap_status', 'mon_is_ibss', 'mon_ap_create_time', 'mon_ap_match_type', 'mon_ap_match_method', 'mon_ap_match_name', 'mon_ap_match_time', 'wms_event_count']
	justip_fields = ['client_ip_address']
	aruba_hosts = []
	aruba_basestations = []

	def get_arguments(cls) :
		"""Returns an array of information used to construct an argumentparser argument."""
		return ['-w', '--wifi', 'store_true', 'Return wireless connection information about the subject (aruba)']

	def startup(self) :
		for requirement in ['host','username','key'] :
			if requirement not in self.config or (requirement in self.config and (self.config[requirement] is '' or type(self.config[requirement]) is not str)):
				self.error.append('Missing required config option ' + requirement)
				return
		self.host = self.config['host']
		self.start_session(self.config['host'],self.config['username'],self.config['key'])
		clientsquery = 'backend-observer-sta-19'
		clientssortfield = 'client_user_name'
		clientsfilter = 'client_conn_type not_equals 0'
		clientsdevicetype = 'sta'
		apsquery = 'backend-observer-mon_bssid-67'
		apssortfield = 'mon_ssid'
		apsfilter = 'mon_ap_status equals 1'
		apsdevicetype = 'mon_bssid'
		self.debug('Retrieving list of Aruba\'s detected basestations...',1)
		self.aruba_basestations = self.perform_list_query(apsquery,self.aps_fields,apssortfield,apsfilter,apsdevicetype)
		self.debug('Retrieving list of Aruba connected clients...',1)
		self.aruba_hosts = self.perform_list_query(clientsquery,self.clients_fields,clientssortfield,clientsfilter,clientsdevicetype)
		self.debug('Searching Aruba information...',1)

	def lookup(self,subject) :
		# We return the first result that matches the subject out of our lists of Aruba information.
		# Because this isn't very effective we need to order our search by order of least-likely to
		# most-likely for the match.

		# Search APs first, since they are less likely to match
		# By IP
		if 'dns' in self.hints and subject in self.hints['dns'] and self.hints['dns'][subject]['addresses'] :
			for ip in self.hints['dns'][subject]['addresses'] :
				for entry in self.aruba_basestations :
					if ip in entry.values() :
						return entry
		# By value
		for entry in self.aruba_basestations :
			lowercase_values = [x.lower() for x in entry.values() if type(x) is str]
			all_values = lowercase_values + [x for x in entry.values() if type(x) is not str]
			for value in all_values :
				if subject.lower() in str(value) :
					return entry
		# Search connected Aruba hosts
		# By IP
		if 'dns' in self.hints and subject in self.hints['dns'] and self.hints['dns'][subject]['addresses'] :
			for ip in self.hints['dns'][subject]['addresses'] :
				for entry in self.aruba_hosts :
					if ip in entry.values() :
						return entry
		# By value
		for entry in self.aruba_hosts :
			lowercase_values = [x.lower() for x in entry.values() if type(x) is str]
			all_values = lowercase_values + [x for x in entry.values() if type(x) is not str]
			for value in all_values :
				if subject.lower() in str(value) :
					return entry
		# Not found
		return None

	def start_session(self,host,username,key) :
		headers = { 'Content-Type': 'application/x-www-form-urlencoded' }
		data = 'opcode=login&uid=' + username + '&passwd=' + key
		self.api_session.post('https://' + host + '/screens/wms/wms.login',verify=False,headers=headers,data=data)
		if 'SESSION' in self.api_session.cookies.get_dict() :
			self.session_key = self.api_session.cookies.get_dict()['SESSION']

	def perform_list_query(self,queryname,columnlist,sortfield,queryfilter,devicetype) :
		""""Performs a list-type XML query against the Aruba UI API"""
		# So uncivilized.
		# Build the basic object for our list query. Takes the parameters that the Aruba controller needs for it's
		# XML api, performs the query, and returns a list of results.
		sortdir = 'asc'
		aruba_queries = Element('aruba_queries')
		query = SubElement(aruba_queries,'query')
		qname = SubElement(query,'qname')
		qname.text = queryname
		type = SubElement(query,'type')
		type.text = 'list'

		list_query = SubElement(query,'list_query')
		device_type = SubElement(list_query,'device_type')
		device_type.text = devicetype
		requested_columns = SubElement(list_query,'requested_columns')
		requested_columns.text = ' '.join(columnlist)
		sort_by_field = SubElement(list_query,'sort_by_field')
		sort_by_field.text = sortfield
		sort_order = SubElement(list_query,'sort_order')
		sort_order.text = sortdir
		pagination = SubElement(list_query,'pagination')
		start_row = SubElement(pagination,'start_row')
		num_rows = SubElement(pagination,'num_rows')

		filter = SubElement(query,'filter')
		global_operator = SubElement(filter,'global_operator')
		global_operator.text = 'and'
		filter_list = SubElement(filter,'filter_list')
		filter_item_entry = SubElement(filter_list,'filter_item_entry')
		field_name = SubElement(filter_item_entry,'field_name')
		field_name.text = queryfilter.split(' ')[0]
		comp_operator = SubElement(filter_item_entry,'comp_operator')
		comp_operator.text = queryfilter.split(' ')[1]
		value = SubElement(filter_item_entry,'value')
		value.text = queryfilter.split(' ')[2]

		# Repeat with page size on the query, aggregating results
		nextstart = 0
		finished = False
		allitems = []
		while not finished :
			start_row.text = str(nextstart)
			num_rows.text = str(self.query_page_size)
			datapayload = ''.join(['query=',tostring(aruba_queries).decode(),'&UIDARUBA=',self.session_key])
			page = self.api_session.post('https://' + self.host + '/screens/cmnutil/execUiQuery.xml',verify=False,timeout=1,data=datapayload)
			self.debug('Got a page of results from Aruba host...',3)
			if page.status_code == 100 or page.status_code == 200 :
				pagecontents = ElementTree.fromstring(page.text)
				rows = pagecontents.iter('row')
				for row in rows:
					if row == {} :
						finished = True
					item = {}
					headers = pagecontents.iter('column_name')
					for value, header in zip(row,headers) :
						item.update({header.text: value.text})
					allitems.append(item)
				# Abort when we get no more rows
				try :
					rows = pagecontents.iter('row')
					next(rows)
				except Exception :
					finished = True
				if page.status_code != 100 and page.status_code !=200 :
					finished = True
				nextstart = nextstart + self.query_page_size
			else :
				finished = True
		self.debug('Retrieved ' + str((len(allitems))) + ' results from query',2)
		return allitems