monolith/core/relay.py

143 lines
4.7 KiB
Python

from twisted.internet.protocol import Protocol, Factory, ClientFactory
from json import dumps, loads
from copy import deepcopy
from datetime import datetime
import main
from utils.logging.log import *
validTypes = ["msg", "notice", "action", "who", "part", "join", "kick", "quit", "nick", "topic", "mode", "conn", "znc", "query", "self", "highlight", "monitor", "err", "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:
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():
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 not i 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 not i in validTypes:
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
def sendRelayNotification(name, cast):
for i in main.relayConnections.keys():
if main.relayConnections[i].authed:
if cast["type"] in main.relayConnections[i].subscriptions:
newCast = deepcopy(cast)
newCast["name"] = name
newCast["time"] = str(datetime.now().isoformat())
main.relayConnections[i].send(dumps(newCast))