monolith/threshold.py

#!/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
from sys import exit

listener = None
connections = {}

def log(data):
    print("[LOG]", data)

def debug(data):
    print("[DEBUG]", data)

class Base(Protocol):
    def __init__(self, addr):
        self.addr = addr
        self.authed = False

    def send(self, data):
        data += "\r\n"
        data = data.encode("utf-8", "replace")
        self.transport.write(data)

    def dataReceived(self, data):
        log("Data received from %s:%s -- %s" % (self.addr.host, self.addr.port, data))

    def connectionMade(self):
        log("Connection from %s:%s" % (self.addr.host, self.addr.port))
        self.send("Hello.")

    def connectionLost(self, reason):
        global connections
        log("Connection lost from %s: %s" % (self.addr.host, reason.getErrorMessage()))
        if not listener == None:
            if self.addr in connections.keys():
                del connections[self.addr]
            else:
                debug("Tried to remove a non-existant connection.")
        else:
            debug("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

def getconfig():
    with open("config.json", "r") as f:
        config = load(f)
        if set(["port", "bind", "usessl"]).issubset(set(config.keys())):
            if config["usessl"] == True:
                if not set(["cert", "key"]).issubset(set(config.keys())):
                    print("SSL is on but certificate or key is not defined")
                    exit(1)
            return config
        else:
            print("Mandatory values missing from config")
            exit(1)

if __name__ == "__main__":
    config = getconfig()

    listener = BaseFactory()
    if config["usessl"] == True:
        reactor.listenSSL(config["port"], listener, DefaultOpenSSLContextFactory(config["key"], config["cert"]), interface=config["bind"])
        log("Threshold running with SSL on %s:%s" % (config["bind"], config["port"]))
    else:
        reactor.listen(config["port"], listener, interface=config["bind"])
        log("Threshold running on %s:%s" % (config["bind"], config["port"]))

    reactor.run()