monolith/core/relay.py

162 lines
4.5 KiB
Python

from json import dumps, loads
from twisted.internet.protocol import Factory, Protocol
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))