from socketserver import StreamRequestHandler, TCPServer, ThreadingMixIn
from struct import pack, unpack
from Crypto.Cipher import AES
from os import urandom
from hashlib import sha256
#from ticket import generate_new_ticket, ticket_valid

#TICKET_MESSAGE = open("ticket_created_message.txt", "r").read()

class ThreadedTCPServer(ThreadingMixIn, TCPServer):
    pass

class ZenCastleServer(StreamRequestHandler):
	def handle(self):
		success, key = self.auth()

		if not success:
			self.write_bool(False)
			self.write_error(b"Auth failed !")
			return

		while True:
			data = self.decrypt(key, self.read_string())

			while not data == b"":
				# Create a ticket
				if data[0] == 1:
					#ticketId = generate_new_ticket()
					self.write_string(self.encrypt(key, "ALLO"))
					data = data[1:]

				# View a ticket status
				elif data[0] == 2:
					ticketId = data[1:17]

					#if ticket_valid(ticketId):
					if True:
						self.write_string(self.encrypt(key, b"Ticket is valid."))
					else:
						self.write_string(self.encrypt(key, b"Ticket is not valid."))

					data = data[17:]

				# Command not found
				else:
					self.write_error(("Command '%d' not found." % data[0]).encode())
					data = b""

	def write_bool(self, value):
		if value:
			self.wfile.write(b"\x01")
		else:
			self.wfile.write(b"\x00")

	def write_error(self, value):
		self.write_string(value)

	def write_string(self, value):
		self.wfile.write(pack("<I", len(value)))
		self.wfile.write(value)

	def read_string(self):
		size = self.rfile.read(4)
		size = unpack("<I", size)[0]

		if size > 128:
			size = 128

		return self.rfile.read(size)

	def get_key_for_username(self, user, challenge):
		with open("accounts.csv", "r") as account_file:
			for account in account_file:
				if account.strip() == "":
					continue

				username, password_hash = account.split(",")
				if username.encode() == user:
					digest = sha256(password_hash.encode())
					digest.update(challenge)
					return digest.digest()

		return None

	def decrypt(self, key, data):
		cipher = AES.new(key, AES.MODE_CFB, b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00')
		return cipher.decrypt(data)

	def encrypt(self, key, data):
		cipher = AES.new(key, AES.MODE_CFB, b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00')
		return cipher.encrypt(data)

	def auth(self):
		client_challenge = self.rfile.read(16)
		server_challenge = urandom(16)

		self.wfile.write(server_challenge)

		username = self.read_string()
		key = self.get_key_for_username(username, client_challenge + server_challenge)

		# User not found !
		if key is None:
			self.write_bool(False)
			return False, None

		auth_challenge = self.rfile.read(16)

		if not self.decrypt(key, auth_challenge) == client_challenge:
			self.write_bool(False)
			return False, None

		self.write_bool(True)
		return True, key

aServer = ThreadedTCPServer(("127.0.0.1", 9090), ZenCastleServer)
aServer.serve_forever()