#!/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()

fileLock = threading.Lock()

commands = {}

def commandlineToList(cmd):
	args = []
	cArg = ""
	escape = False
	quoted = False
	for letter in cmd:
		if escape == True:
			cArg += letter
			escape = False
			continue
		
		if letter == "\\":
			escape = True
			continue
		
		#if quoted == False and letter == ",":
		if letter == ",":
			if cArg == "": continue
			args.append(cArg)
			cArg = ""
			continue
		
		#if letter == '"':
		#	quoted = not quoted
		#	continue
		
		cArg += letter
	
	args.append(cArg)
	
	return args

def listToCommandline(lst):
	cmd = ""
	for arg in lst:
		arg = arg.replace("\\","\\\\")
		arg = arg.replace(",","\\,")
		#arg = arg.replace('"','\\"')
		#if " " in arg: arg = '"' +arg+ '"'
		cmd += arg + ","
	
	return cmd[:-1]

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

def runCode(str, lcs = False, description = "loose-code"):
	if lcs == False: lcs = {}
	code = compile(str,description,"exec")
	exec(code,globals(),lcs)
	return lcs
	
def runScript(sf, lcs = False):
	if lcs == False: lcs = {}
	with open(sf) as script:
		runCode(script.read(),lcs,sf)
		return lcs

def getModlist(path):
	modList = []
	for root,dirs,files in os.walk(path):
		for file in dirs:
			ffile = p(root,file)
			lfile = ffile.replace(path + os.path.sep,"",1)
			if lfile[0] == "-": continue
			if lfile[0] == "[" and lfile[-1] == "]":
				modList = modList + sorted(getModlist(ffile))
				continue
			
			modList.append(ffile)
		break
	
	return modList

modulesLoaded = []
modulePath = p(sp,"modules")
def moduleRun(localModule):
	if not localModule in modulesLoaded: modulesLoaded.append(localModule)
	print("> " +localModule+ "...")
	runScript(p(modulePath,localModule,"module.py"))

def moduleDepends(localModules):
	if type(localModules) == str: localModules = [localModules]
	
	for localModule in localModules:
		if localModule in modulesLoaded: return
		print("depend ",end="")
		moduleRun(localModule)

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") + b"\x00" + data)

def getResponse(connection):
	data = b''
	data = connection.recv(4)
	if not data: return False
	nul = connection.recv(1)
	if not nul: return False
	if nul != b"\x00": 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 runCommand(self,cmd):
		command = False
		if not cmd[0] in commands:
			return ["error","nonfatal","command_not_found"]		
		command = commands[cmd[0]]
		rtn = command["function"](self,cmd)
		
		sleep = pauseBetweenCommands
		if "sleep" in command:
			sleep = command["sleep"]
		if sleep > 0: time.sleep(sleep)
		
		return rtn
	
	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
				
				queue = False
				with connectionsLock:
					queue = connections[self.connectionId]["threadOut"].queue
					if queue.qsize() >= maxQueueSize:
						closeConnection(self.connectionId)
						return
				
				dataString = data.decode(encoding="utf-8")
				commandList = commandlineToList(dataString)
				queue.put(listToCommandline(self.runCommand(commandList)).encode(encoding="utf-8"))
		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():
	print("Loading modules...")
	for path in getModlist(modulePath):
		if os.path.isfile(p(path,"module.py")):
			localModule = path.replace(modulePath + os.path.sep,"",1)
			if not localModule in modulesLoaded:
				moduleRun(localModule)
	
	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()