168 lines
4.6 KiB
Python
168 lines
4.6 KiB
Python
from copy import deepcopy
|
|
from json import dumps, loads
|
|
|
|
from twisted.internet.protocol import ClientFactory, Factory, Protocol
|
|
|
|
import main
|
|
from utils.logging.log import *
|
|
|
|
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:
|
|
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("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))
|