from twisted.internet.protocol import Protocol, Factory from json import dumps, loads import main from utils.logging.log import log, warn validTypes = [ "msg", "notice", "action", "who", "part", "join", "kick", "quit", "nick", "topic", "mode", "conn", "znc", "query", "self", "highlight", ] 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: # noqa: E722 self.sendErr("MALFORMED") return if "type" not 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(): if self.authed: self.handleSubscribe(parsed["subscribe"]) return else: self.sendErr("DENIED") return elif "unsubscribe" in parsed.keys(): if self.authed: self.handleUnsubscribe(parsed["unsubscribe"]) return else: self.sendErr("DENIED") 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 i not in validTypes: 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 i not in validTypes: self.sendErr("NONEXISTANT") return if i not 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("Relay 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 def sendRelayNotification(cast): for i in main.relayConnections.keys(): if main.relayConnections[i].authed: if cast["type"] in main.relayConnections[i].subscriptions: main.relayConnections[i].send(dumps(cast))