#!/usr/bin/env python3
import sys

stopOnException = True
oldexcepthook = sys.excepthook
def newexcepthook(type,value,traceback):
	oldexcepthook(type,value,traceback)
	if stopOnException: 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)

import socket
import threading
import queue
import traceback

# SETTINGS
serverAddr = ("127.0.0.1",1337)
moduleDir = p(sp,"modules")
heartbeatTime = 300
maxRequestSize = 4096
# SETTINGS END

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

handlers = {}
handlers["modulesLoaded"] = []
handlers["preConnect"] = []
handlers["connect"] = []
handlers["preCommand"] = []
handlers["command"] = []

commands = {}

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 callHandler(handlerName,env = {},*args,**kwargs):
	if handlerName in handlers:
		for handlerFunc in handlers[handlerName]:
			if handlerFunc(env,*args,**kwargs):
				return True
	
	return False

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: return False
	return connection.recv(requestLength)

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

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 letter == ",":
			if cArg == "": continue
			args.append(cArg)
			cArg = ""
			continue
		
		cArg += letter
	
	args.append(cArg)
	
	return args

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

def runCommand(self,command,*args):
	callHandler("preCommand",locals())
	if not command in commands:
		rtn = ["error","nonfatal","invalid_command","Command does not exist"]
		callHandler("command",locals())
		return rtn
	
	rtn = commands[command](self,command,*args)
	if not rtn: rtn = ["OK"]
	callHandler("command",locals())
	return rtn

class connectionThreadIn(threading.Thread):
	def __init__(self,cid,connection,address):
		threading.Thread.__init__(self)
		self.cid = cid
		self.connection = connection
		self.address = address
	
	def routine(self):
		while True:
			data = getResponse(self.connection)
			if data == False: return
			commandList = commandlineToList(data.decode("utf-8"))
			rtn = runCommand(self,*commandList)
			with connectionsLock:
				connections[self.cid]["threadOut"].queue.put(listToCommandline(rtn))
	
	def run(self):
		try:
			self.routine()
		except:
			print(traceback.format_exc())
		
		try:
			self.connection.close()
		except:
			pass
		
		with connectionsLock:
			try:
				connections[self.cid]["threadOut"].queue.put(False)
			except:
				pass
			
			try:
				del connections[self.cid]
			except:
				pass

class connectionThreadOut(threading.Thread):
	def __init__(self,cid,connection,address):
		threading.Thread.__init__(self)
		self.cid = cid
		self.connection = connection
		self.address = address
		self.queue = queue.Queue()
	
	def routine(self):
		while True:
			data = self.queue.get(timeout=heartbeatTime)
			if data == False:
				return
			sendResponse(self.connection,data.encode("utf-8"))
	
	def run(self):
		try:
			self.routine()
		except:
			print(traceback.format_exc())
		
		try:
			self.connection.close()
		except:
			pass
		
		try:
			with connectionsLock:
				del connections[self.cid]
		except:
			pass

def main():
	if os.path.isfile("modules.txt"):
		print("Loading modules...")
		with open("modules.txt","r") as modulesFile:
			for line in modulesFile:
				line = line.split("#",1)[0].strip(" \t\r\n")
				if line == "": continue
				print("> " +line+ " ...")
				moduleFile = p(moduleDir,line)
				runScript(moduleFile,locals())
		print("OK.\n")
	
	callHandler("modulesLoaded",locals())
	
	global connectionId
	global serverSocket
	serverSocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
	serverSocket.bind(serverAddr)
	serverSocket.listen(65535)
	print("Serving at " +str(serverAddr[0])+ ":" +str(serverAddr[1])+ ".")
	
	while True:
		connection = False
		try:
			connection,address = serverSocket.accept()
			connection.settimeout(heartbeatTime)
			
			if callHandler("preConnect",locals()):
				connection.close()
				continue
			
			with connectionsLock:
				connectionId += 1
				threadIn = connectionThreadIn(str(connectionId),connection,address)
				threadOut = connectionThreadOut(str(connectionId),connection,address)
				connections[str(connectionId)] = {
					"connection":	connection,
					"address":		address,
					"threadIn":		threadIn,
					"threadOut":	threadOut
				}
				
				threadIn.start()
				threadOut.start()
			
			callHandler("connect",locals())
		except:
			print(traceback.format_exc())
			try:
				connection.close()
			except:
				pass

if __name__ == '__main__':
	main()