monolith/threshold

226 lines
6.9 KiB
Python
Executable File

#!/usr/bin/env python
from twisted.internet import reactor
from twisted.internet.ssl import DefaultOpenSSLContextFactory
from twisted.internet.protocol import Protocol, Factory
from json import load, dump, loads
from sys import exit
listener = None
connections = {}
def log(data):
print("[LOG]", data)
def debug(data):
print("[DEBUG]", data)
def warn(data):
print("[WARNING]", data)
def error(data):
print("[ERROR]", data)
exit(1)
def sendData(addr, data):
connections[addr].send(data)
def sendSuccess(addr, data):
sendData(addr, "[y] " + data)
def sendFailure(addr, data):
sendData(addr, "[n] " + data)
def sendInfo(addr, data):
sendData(addr, "[i] " + data)
class Base(Protocol):
def __init__(self, addr):
self.addr = addr
self.authed = False
if config["UsePassword"] == False:
self.authed = True
def send(self, data):
data += "\r\n"
data = data.encode("utf-8", "replace")
self.transport.write(data)
def dataReceived(self, data):
data = data.decode("utf-8", "replace")
log("Data received from %s:%s -- %s" % (self.addr.host, self.addr.port, repr(data)))
helper.parseCommand(self.addr, self.authed, data)
def connectionMade(self):
log("Connection from %s:%s" % (self.addr.host, self.addr.port))
self.send("Hello.")
def connectionLost(self, reason):
global connections
self.authed = False
log("Connection lost from %s:%s -- %s" % (self.addr.host, self.addr.port, reason.getErrorMessage()))
if not listener == None:
if self.addr in connections.keys():
del connections[self.addr]
else:
warn("Tried to remove a non-existant connection.")
else:
warn("Tried to remove a connection from a listener that wasn't running.")
class BaseFactory(Factory):
def buildProtocol(self, addr):
global connections
entry = Base(addr)
connections[addr] = entry
return entry
def send(self, addr, data):
global connections
if addr in connections.keys():
connection = connections[addr]
connection.send(data)
else:
return
class Helper(object):
def getConfig(self):
with open("config.json", "r") as f:
config = load(f)
if set(["Port", "BindAddress", "UseSSL", "UsePassword"]).issubset(set(config.keys())):
if config["UseSSL"] == True:
if not set(["ListenerKey", "ListenerCertificate"]).issubset(set(config.keys())):
error("SSL is on but certificate or key is not defined")
if config["UsePassword"] == True:
if not "Password" in config.keys():
error("Password authentication is on but password is not defined")
return config
else:
error("Mandatory values missing from config")
def getPool(self):
with open("pool.json", "r") as f:
data = f.read()
if not data == "":
pool = loads(data.strip())
return pool
else:
return {}
def savePool(self):
global pool
with open("pool.json", "w") as f:
dump(pool, f, indent=4)
return
def getHelp(self):
with open("help.json", "r") as f:
help = load(f)
return help
def incorrectUsage(self, addr, mode):
if mode == None:
sendFailure(addr, "Incorrect usage")
return
if mode in help.keys():
sendFailure(addr, "Usage: " + help[mode])
return
def parseCommand(self, addr, authed, data):
global pool
data = data.strip()
spl = data.split()
obj = connections[addr]
success = lambda data: sendSuccess(addr, data)
failure = lambda data: sendFailure(addr, data)
info = lambda data: sendInfo(addr, data)
incUsage = lambda mode: self.incorrectUsage(addr, mode)
length = len(spl)
if len(spl) > 0:
cmd = spl[0]
else:
failure("No text was sent")
return
if authed == True:
if cmd == "help":
helpMap = []
for i in help.keys():
helpMap.append("%s: %s" % (i, help[i]))
info("\n".join(helpMap))
return
elif cmd == "pass":
info("You are already authenticated")
return
elif cmd == "logout":
obj.authed = False
success("Logged out")
return
elif cmd == "list":
poolMap = []
for i in pool.keys():
poolMap.append("%s: %s" % (i, pool[i]))
info("\n".join(poolMap))
return
elif cmd == "connect":
if length == 6:
if spl[1] in pool.keys():
failure("Name already exists: %s" % spl[1])
return
protocol = spl[4].lower()
if not protocol in ["ssl", "plain"]:
incUsage("connect")
return
pool[spl[1]] = { "address": spl[2],
"port": spl[3],
"protocol": protocol,
"nickname": spl[5],
}
success("Successfully created bot")
self.savePool()
return
else:
incUsage("connect")
return
else:
incUsage(None)
return
else:
if cmd == "pass" and length == 2:
if spl[1] == config["Password"]:
success("Authenticated successfully")
obj.authed = True
return
else:
failure("Password incorrect")
obj.transport.loseConnection()
return
else:
incUsage(None)
return
if __name__ == "__main__":
helper = Helper()
config = helper.getConfig()
pool = helper.getPool()
help = helper.getHelp()
listener = BaseFactory()
if config["UseSSL"] == True:
reactor.listenSSL(config["Port"], listener, DefaultOpenSSLContextFactory(config["ListenerKey"], config["ListenerCertificate"]), interface=config["BindAddress"])
log("Threshold running with SSL on %s:%s" % (config["BindAddress"], config["Port"]))
else:
reactor.listenTCP(config["Port"], listener, interface=config["BindAddress"])
log("Threshold running on %s:%s" % (config["BindAddress"], config["Port"]))
reactor.run()