#!/usr/bin/python3
import asyncio
import websockets
import time
import hashlib
import sys
import settings
import utils
VERSION = "baitv-daemon v0.2.4"
USERS = set()
current_time = lambda: int(round(time.time() * 1000))
timestamp = lambda: int(round(time.time()))
if len(sys.argv) > 1:
settings.port = sys.argv[1]
def update_state(k, v):
state = utils.load_state(settings.state_file)
state[k] = v
utils.save_state(state, settings.state_file)
class Client:
def __init__(self, ws, ip):
self.ws = ws
self.ip = ip
self.last_chat = 0
self.mod = None
self.cmds = {
"msg": (self.on_msg, 2, False),
"notice": (self.on_notice, 2, True),
"login": (self.on_login, 1, False),
"logout": (self.on_logout, 0, True),
"kick": (self.on_kick, 1, True),
"ban": (self.on_ban, 1, True),
"settitle": (self.on_settitle, 1, True),
"setsource": (self.on_setsource, 1, True),
}
if self.ip == "127.0.0.1":
self.mod = "local"
else:
self.ip_hash = hashlib.md5(self.ip.encode('ascii')).hexdigest()[:5]
print("IP: {} ({})".format(self.ip, self.ip_hash))
@asyncio.coroutine
def broadcast(self, msg, exclude=[]):
users_to_send = [user for user in USERS if user not in exclude]
if users_to_send:
yield from asyncio.wait([user.ws.send(msg) for user in users_to_send])
@asyncio.coroutine
def on_msg(self, args):
if current_time() - self.last_chat > settings.flood_time:
color, msg = args
for user in USERS:
if user != self:
if user.mod:
yield from user.ws.send("FMSG:{}:{}:{}".format(self.ip_hash, color, msg))
else:
yield from user.ws.send("MSG:{}:{}".format(color, msg))
self.last_chat = current_time()
else:
yield from self.ws.send("FLOOD")
@asyncio.coroutine
def on_notice(self, args):
(showname, msg) = args
if showname == "1":
yield from self.broadcast("NOTICE:{}:{}".format(self.mod, msg), [self])
else:
yield from self.broadcast("NOTICE::{}".format(msg), [self])
@asyncio.coroutine
def on_login(self, args):
valid_psks = utils.parse_streamers_file(settings.streamers_file)
psk = args[0]
if psk in valid_psks:
print("User logged in")
self.mod = valid_psks[psk]
yield from self.ws.send("LOGIN_OK:{}".format(self.mod))
else:
yield from self.ws.send("BADPASS")
@asyncio.coroutine
def on_logout(self, args):
self.mod = None
yield from self.ws.send("LOGOUT_OK")
@asyncio.coroutine
def on_kick(self, args):
users_to_kick = [user for user in USERS if args[0] == user.ip_hash]
if users_to_kick:
yield from self.broadcast("KICKED", users_to_kick)
for user in users_to_kick:
yield from user.ws.send("YOU_KICKED")
yield from user.ws.close(4000, "Has sido expulsado.")
yield from self.ws.send("KICK_OK")
print("Kicked " + args[0])
@asyncio.coroutine
def on_ban(self, args):
users_to_ban = [user for user in USERS if args[0] == user.ip_hash]
if users_to_ban:
yield from self.broadcast("BANNED", users_to_ban)
utils.add_ban(users_to_ban[0].ip, settings.bans_file)
for user in users_to_ban:
yield from user.ws.send("YOU_BANNED")
yield from user.ws.close(4001, "Has sido baneado.")
yield from self.ws.send("BAN_OK")
print("Banned " + ip_to_ban)
@asyncio.coroutine
def on_settitle(self, args):
update_state('title', args[0])
yield from self.broadcast("TITLE:{}".format(args[0]), [self])
@asyncio.coroutine
def on_setsource(self, args):
t = timestamp()
update_state('source', args[0])
update_state('source_time', t)
yield from self.broadcast("SOURCE:{}:{}".format(args[0], t), [self])
@asyncio.coroutine
def handle_message(self, msg_in):
if len(msg_in) > settings.max_cmd_length:
yield from self.ws.send("FUCKOFF")
return
args = msg_in.split(":")
cmd = args[0].lower()
if cmd in self.cmds:
cmd_to_run, required_args, need_mod = self.cmds[cmd]
if len(args) - 1 == required_args:
if need_mod:
if self.mod:
yield from cmd_to_run(args[1:])
else:
yield from self.ws.send("FORBIDDEN")
else:
yield from cmd_to_run(args[1:])
else:
yield from self.ws.send("INVALID")
else:
yield from self.ws.send("WHAT")
@asyncio.coroutine
def broadcast(msg):
print("broadcasting " + msg)
if USERS:
yield from asyncio.wait([user.ws.send(msg) for user in USERS])
@asyncio.coroutine
def notify_count():
total_users = len(USERS)
yield from broadcast("COUNT:{}".format(total_users))
@asyncio.coroutine
def register(websocket):
USERS.add(websocket)
yield from notify_count()
@asyncio.coroutine
def unregister(websocket):
USERS.remove(websocket)
yield from notify_count()
@asyncio.coroutine
def baitv_daemon(websocket, path):
print("New client ({})".format(websocket.remote_address))
real_ip = websocket.request_headers['X-Real-IP']
if not real_ip:
real_ip = websocket.remote_address[0]
this_client = Client(websocket, real_ip)
yield from websocket.send("WELCOME:{}".format(VERSION))
#print(list(websocket.request_headers.raw_items()))
if utils.is_banned(real_ip, settings.bans_file):
yield from websocket.send("YOU_BANNED")
yield from websocket.close(4002, "Estás baneado.")
return
if real_ip != "127.0.0.1":
yield from register(this_client)
try:
while True:
msg_in = yield from websocket.recv()
print("< {}".format(msg_in))
yield from this_client.handle_message(msg_in)
except websockets.ConnectionClosed as e:
print("Connection was closed. ({}|{})".format(e.code, e.reason))
finally:
print("Client disconnected ({})".format(websocket.remote_address))
if real_ip != "127.0.0.1":
yield from unregister(this_client)
print("baitv-daemon {}".format(VERSION))
if settings.use_ssl:
import ssl
print("Loading certificates...")
ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS)
ssl_context.load_cert_chain(settings.cert, keyfile=settings.cert_key)
print("Starting secure server on {}:{}...".format(settings.addr, settings.port))
start_server = websockets.serve(baitv_daemon, settings.addr, settings.port, ssl=ssl_context)
else:
print("Starting plain server on {}:{}...".format(settings.addr, settings.port))
start_server = websockets.serve(baitv_daemon, settings.addr, settings.port)
asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()
