import asyncio
import traceback
import logging
import ssl



from asn1crypto.x509 import Certificate

class SSLWrapper:
	def __init__(self):
		self.comms = None
		self.transport_in = None
		self.transport_out = None
		self.ssl_ctx = None
		self.transport = None
		self.ssl_in_buff = None
		self.ssl_out_buff = None
		self.ssl_obj = None
		self.handshake_done_evt = None
		self.handshake_timeout = 10
		self.__ssl_in_q = None
		self.stop_evt = None
		self.server_certificate = None

		self.__enc_task = None
		self.__dec_task = None
		self.__process_task = None
	
	async def disconnect(self):
		try:
			if self.__enc_task is not None:
				self.__enc_task.cancel()
			if self.__dec_task is not None:
				self.__dec_task.cancel()
			if self.__process_task is not None:
				self.__process_task.cancel()

			if self.transport is not None:
				self.transport_out.put_nowait(b'')
				await self.transport.disconnect()
		except Exception as e:
			pass
	
	async def get_server_pubkey(self):
		try:
			if self.server_certificate is None:
				return None
			cert = Certificate.load(self.server_certificate)
			return cert['tbs_certificate']['subject_public_key_info']['public_key'].dump()[5:] #why?
		except Exception as e:
			return e
	
	def get_server_cert(self):
		if self.server_certificate is not None:
			return self.server_certificate
	
	async def setup(self):
		try:
			self.handshake_done_evt = asyncio.Event()
			self.stop_evt = asyncio.Event()
			self.__ssl_in_q = asyncio.Queue()
			self.ssl_in_buff = ssl.MemoryBIO()
			self.ssl_out_buff = ssl.MemoryBIO()
			self.ssl_obj = self.ssl_ctx.wrap_bio(self.ssl_in_buff, self.ssl_out_buff, server_side=False, server_hostname = None)
			_, err = await self.__do_handshake()
			if err is not None:
				raise err
			self.server_certificate = self.ssl_obj.getpeercert(binary_form=True)
			self.__enc_task = asyncio.create_task(self.__encrypt_ssl_rec())
			self.__dec_task = asyncio.create_task(self.__decrypt_ssl_rec())
			self.__process_task = asyncio.create_task(self.__read_ssl_record())
			return True, None

		except Exception as e:
			return None, e
	
	async def __encrypt_ssl_rec(self):
		try:
			while True:
				try:
					raw_data = await self.comms.in_queue.get()
					if raw_data == b'':
						await self.comms.in_queue.put(None)
						return
					
					self.ssl_obj.write(raw_data)
					while True:
						try:
							enc_data = self.ssl_out_buff.read()
							if enc_data == b'':
								break
							await self.transport_out.put(enc_data)
						except Exception as e:
							raise e
				except Exception as e:
							raise e

		except asyncio.CancelledError:
			return
		except Exception as e:
			logging.debug('__encrypt_ssl_rec %s' % e)
	
	async def __decrypt_ssl_rec(self):
		"""
		Decrypting the encrypted SSL records from the internal queue and putting them in the queue for the uppr layers to use
		"""
		try:
			while not self.stop_evt.is_set():
				ssl_data = await self.__ssl_in_q.get()
				if ssl_data == b'':
					await self.comms.out_queue.put((ssl_data, None))
					#print('Connection terminated')
					return
				self.ssl_in_buff.write(ssl_data)

				data_buff = b''
				while not self.stop_evt.is_set():
					try:
						data_buff += self.ssl_obj.read()
					except ssl.SSLWantReadError:
						break
					except Exception as e:
						raise e
				#print('data_buff %s' % data_buff)
				if data_buff != b'':
					await self.comms.out_queue.put((data_buff, None))

		except asyncio.CancelledError:
			return
		except Exception as e:
			logging.debug('__decrypt_ssl_rec %s' % e)

	async def __read_ssl_record(self):
		try:
			"""
			buffering the incoming SSL (encrypted) recodrs and putting them in an internal queue one record at a time
			"""
			buffer = b''
			length = None
			while not self.stop_evt.is_set():
				if length is None and len(buffer) >= 6:
					length = int.from_bytes(buffer[3:5], byteorder = 'big', signed = False)
				
				if length is not None and len(buffer) >= length + 5:
					#print('LB raw %s' % len(buffer[:length+5]))
					await self.__ssl_in_q.put(buffer[:length+5])
					buffer = buffer[length+5:]
					length = None
					continue
				
				data, err = await self.transport_in.get()
				if err is not None:
					raise err
				if data == b'':
					await self.__ssl_in_q.put(b'')
					return
				buffer+= data
				

		except asyncio.CancelledError:
			return

		except Exception as e:
			logging.debug('__read_ssl_record %s' % e)
			await self.__ssl_in_q.put(b'')

		#finally:
		#	self.stop_evt.set()
	
	async def __do_handshake(self):
		try:
			ctr = 0
			while not self.stop_evt.is_set():
				ctr += 1
				#print('DST Performing handshake!')
				try:
					self.ssl_obj.do_handshake()
				except ssl.SSLWantReadError:
					#print('DST want %s' % ctr)
					while True:
						client_hello = self.ssl_out_buff.read()
						if client_hello != b'':
							#print('DST client_hello %s' % len(client_hello))
							await self.transport_out.put(client_hello)
						else:
							break
					
					#print('DST wating server hello %s' % ctr)
					server_hello, err = await self.transport_in.get()
					if err is not None:
						raise err
					if server_hello == b'':
						raise Exception('Server closed the connection!')
					#print('DST server_hello %s' % len(server_hello))
					self.ssl_in_buff.write(server_hello)

					continue
				except:
					raise
				else:
					#print('DST handshake ok %s' % ctr)
					#server_fin = tls_out_buff.read()
					#print('DST server_fin %s ' %  server_fin)
					#await out_q.put(server_fin)
					break			
			
			self.handshake_done_evt.set()
			return True, None

		except Exception as e:
			logging.debug('__do_handshake %s' % e)
			return None, e

	@staticmethod
	async def from_commsettings(comms, transport_in, transport_out):
		try:
			tunnel = SSLWrapper()
			tunnel.transport_in = transport_in
			tunnel.transport_out = transport_out
			tunnel.comms = comms
			tunnel.handshake_timeout = comms.handshake_timeout
			if comms.ssl_ctx is not None:
				tunnel.ssl_ctx = comms.ssl_ctx
			#elif comms.ssl_cert is not None:
			#	raise NotImplementedError()
			else:
				tunnel.ssl_ctx = ssl.create_default_context()
				tunnel.ssl_ctx.check_hostname = False
				tunnel.ssl_ctx.verify_mode = ssl.CERT_NONE
			_, err = await tunnel.setup()
			if err is not None:
				raise err

			return tunnel, None
		except Exception as e:
			return None, e