Implement the relay channel and command for generating tokens

core/relay.py

@@ -0,0 +1,120 @@
+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