#!/usr/bin/env python3
import sys

oldexcepthook = sys.excepthook
def newexcepthook(type,value,traceback):
	oldexcepthook(type,value,traceback)
	input("Press ENTER to quit.")
sys.excepthook = newexcepthook

import os
p = os.path.join
pUp = os.path.dirname
s = False
if getattr(sys, 'frozen', False) and hasattr(sys, '_MEIPASS'):
	s = os.path.realpath(sys.executable)
else:
	s = os.path.realpath(__file__)
sp = pUp(s)

# script start
import threading
import queue
import socket
import traceback
import time
import colorama
colorama.init()

maxConnections = 10000
maxConnectionsPerIp = 10
maxQueueSize = 1000
maxRequestSize = 4096
pauseBetweenCommands = 0.1

serverAddr = ("127.0.0.1",21779)
serverSocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

connectionsLock = threading.Lock()
connections = {}
connectionsId = 0

heartbeatTime = 600

threadCount = 0
threadCountLock = threading.Lock()

printLock = threading.Lock()
def tprint(st):
	with printLock:
		print(st)

def addThread():
	global threadCount
	with threadCountLock:
		threadCount += 1
		tprint(colorama.Fore.YELLOW + colorama.Style.BRIGHT + "Thread opened. Threads: " +str(threadCount)+ " (Actual: " +str(threading.active_count())+ ")" + colorama.Style.RESET_ALL)

def removeThread():
	global threadCount
	with threadCountLock:
		threadCount -= 1
		tprint(colorama.Fore.YELLOW + colorama.Style.BRIGHT + "Thread closed. Threads: " +str(threadCount)+ " (Actual: " +str(threading.active_count())+ ")" + colorama.Style.RESET_ALL)

def sendResponse(connection,data):
	connection.sendall(len(data).to_bytes(4,"big") + data)

def getResponse(connection):
	data = b''
	data = connection.recv(4)
	if not data: return False
	requestLength = int.from_bytes(data,"big")
	if requestLength > maxRequestSize: raise Exception("security","request_too_large")
	return connection.recv(requestLength)

def closeConnection(connectionId):
	if not connectionId in connections: return False
	try:
		connections[connectionId]["connection"].close()
	except Exception as e:
		tprint("Failed to close connection: " +str(e))
		pass
	
	try:
		connections[connectionId]["threadOut"].queue.put(False)
	except:
		with printLock:
			print(colorama.Fore.GREEN + colorama.Style.BRIGHT)
			traceback.print_exc()
			print(colorama.Style.RESET_ALL)
	
	del connections[connectionId]
	return True

class connectionThreadOut(threading.Thread):
	def __init__(self,connectionId):
		threading.Thread.__init__(self)
		self.queue = queue.Queue()
		self.connectionId = connectionId
	
	def getConnection(self):
		with connectionsLock:
			if self.connectionId in connections:
				return connections[self.connectionId]["connection"]
			return False
	
	def run(self):
		try:
			while True:
				data = self.queue.get(timeout=heartbeatTime)
				if data == False: return
				
				connection = self.getConnection()
				if not connection:
					with connectionsLock: closeConnection(self.connectionId)
					return
				
				sendResponse(connection,data)
		except Exception as e:
			with connectionsLock: closeConnection(self.connectionId)
			with printLock:
				print(colorama.Fore.GREEN + colorama.Style.BRIGHT)
				traceback.print_exc()
				print(colorama.Style.RESET_ALL)
		finally:
			removeThread()

class connectionThreadIn(threading.Thread):
	def __init__(self,connectionId):
		threading.Thread.__init__(self)
		self.connectionId = connectionId
	
	def getConnection(self):
		with connectionsLock:
			if self.connectionId in connections:
				return connections[self.connectionId]["connection"]
			return False
	
	def run(self):
		try:
			while True:
				connection = self.getConnection()
				if not connection:
					with connectionsLock: closeConnection(self.connectionId)
					return
				
				data = getResponse(connection)
				if data == False:
					with connectionsLock: closeConnection(self.connectionId)
					return
				
				with connectionsLock:
					if self.connectionId in connections:
						queue = connections[self.connectionId]["threadOut"].queue
						if queue.qsize() >= maxQueueSize:
							closeConnection(self.connectionId)
							return
						queue.put(data)
						time.sleep(pauseBetweenCommands)
		except Exception as e:
			with connectionsLock: closeConnection(self.connectionId)
			with printLock:
				print(colorama.Fore.GREEN + colorama.Style.BRIGHT)
				traceback.print_exc()
				print(colorama.Style.RESET_ALL)
		finally:
			removeThread()

def main():
	global connectionsId
	serverSocket.bind(serverAddr)
	serverSocket.listen(65535)
	
	while True:
		connection,address = serverSocket.accept()
		connection.settimeout(heartbeatTime)
		
		with connectionsLock:
			# Count connections
			connectionsCount = 0
			connectionsCountIp = 0
			for connectionId in connections:
				connectionsCount += 1
				if connections[connectionId]["address"][0] == address[0]:
					connectionsCountIp += 1
			
			if connectionsCount >= maxConnections:
				tprint("Connection closed - too many clients.")
				closeConnection(connectionId)
				continue
			
			if connectionsCountIp >= maxConnectionsPerIp:
				tprint("Connection closed - same IP connected too many times.")
				closeConnection(connectionId)
				continue
			
			# Create connection
			connectionsId += 1
			threadIn = connectionThreadIn(str(connectionsId))
			threadOut = connectionThreadOut(str(connectionsId))
			connections[str(connectionsId)] = {
				"connection":	connection,
				"address":		address,
				"threadOut":	threadOut,
				"threadIn":		threadIn,
				"user":			False
			}
			
			addThread()
			addThread()
			threadOut.start()
			threadIn.start()

if __name__ == '__main__':
	main()