monolith/core/relay.py

from twisted.internet.protocol import Protocol, Factory, ClientFactory
from json import dumps, loads
import main
from utils.logging.log import *

class Relay(Protocol):
    def __init__(self, addr):
        self.addr = addr
        self.authed = False
        self.subscriptions = []

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

    def sendErr(self, data):
        self.send(dumps({"type": "error", "reason": data}))
        return

    def sendMsg(self, data):
        self.send(dumps(data))

    def dataReceived(self, data):
        data = data.decode("utf-8", "replace")
        try:
            parsed = loads(data)
        except:
            self.sendErr("MALFORMED")
            return
        if not "type" in parsed.keys():
            self.sendErr("NOTYPE")
            return
        if parsed["type"] == "hello":
            if set(["key", "hello"]).issubset(set(parsed)):
                self.handleHello(parsed)
            else:
                self.sendErr("WRONGFIELDS")
                return
            return
        elif parsed["type"] == "control":
            if "subscribe" in parsed.keys():
                self.handleSubscribe(parsed["subscribe"])
                return
            elif "unsubscribe" in parsed.keys():
                self.handleUnsubscribe(parsed["unsubscribe"])
                return
            else:
                self.sendErr("UNCLEAR")
                return
        else:
            self.sendErr("UNCLEAR")
            return

    def handleSubscribe(self, lst):
        if not isinstance(lst, list):
            self.sendErr("NOTLIST")
            return
        for i in lst:
            if not i in ["msg", "notice", "action", "who", "part", "join", "kick", "quit", "nick", "topic", "mode"]:
                self.sendErr("NONEXISTANT")
                return
            if i in self.subscriptions:
                self.sendErr("SUBSCRIBED")
                return
            self.subscriptions.append(i)
        self.sendMsg({"type": "success"})
        return

    def handleUnubscribe(self, lst):
        if not isinstance(lst, list):
            self.sendErr("NOTLIST")
            return
        for i in lst:
            if not i in ["msg", "notice", "action", "who", "part", "join", "kick", "quit", "nick", "topic", "mode"]:
                self.sendErr("NONEXISTANT")
                return
            if not i in self.subscriptions:
                self.sendErr("NOTSUBSCRIBED")
                return
            del self.subscriptions[i]
        self.sendMsg({"type": "success"})
        return

    def handleHello(self, parsed):
        if parsed["key"] in main.tokens.keys():
            if parsed["hello"] == main.tokens[parsed["key"]]["hello"] and main.tokens[parsed["key"]]["usage"] == "relay":
                self.sendMsg({"type": "hello", "hello": main.tokens[parsed["key"]]["counter"]})
                self.authed = True
            else:
                self.transport.loseConnection()
                return
        else:
            self.sendErr("NOKEY")
            return

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

    def connectionLost(self, reason):
        self.authed = False
        log("Relay connection lost from %s:%s -- %s" % (self.addr.host, self.addr.port, reason.getErrorMessage()))
        if self.addr in main.relayConnections.keys():
            del main.relayConnections[self.addr]
        else:
            warn("Tried to remove a non-existant relay connection.")

class RelayFactory(Factory):
    def buildProtocol(self, addr):
        entry = Relay(addr)
        main.relayConnections[addr] = entry
        return entry

    def send(self, addr, data):
        if addr in main.relayConnections.keys():
            connection = main.relayConnections[addr]
            connection.send(data)
        else:
            return