#!/usr/bin/python3 import asyncio import websockets import time import hashlib import sys import settings import utils VERSION = "baitv-daemon v0.2.5" USERS = set() HISTORY = [] 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), "history":(self.on_history, 1, False), "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)) async def broadcast(self, msg, exclude=[]): users_to_send = [user for user in USERS if user not in exclude] if users_to_send: _ = [await user.ws.send(msg) for user in users_to_send] async def on_msg(self, args): t = current_time() if t - self.last_chat > settings.flood_time: color, msg = args for user in USERS: if user != self: if user.mod: await user.ws.send("FMSG:{}:{}:{}".format(self.ip_hash, color, msg)) else: await user.ws.send("MSG:{}:{}".format(color, msg)) self.last_chat = t HISTORY.append((t, color, msg)) if len(HISTORY) > settings.history_length: HISTORY.pop(0) else: await self.ws.send("FLOOD") async def on_notice(self, args): (showname, msg) = args if showname == "1": await self.broadcast("NOTICE:{}:{}".format(self.mod, msg), [self]) else: await self.broadcast("NOTICE::{}".format(msg), [self]) async def on_history(self, args): if HISTORY: msgs = min(len(HISTORY), int(args[0])) * -1 for msg in HISTORY[msgs:]: await self.ws.send("HMSG:{}:{}:{}".format(*msg)) await self.ws.send("HISTORY_OK") async 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] await self.ws.send("LOGIN_OK:{}".format(self.mod)) else: await self.ws.send("BADPASS") async def on_logout(self, args): self.mod = None await self.ws.send("LOGOUT_OK") async def on_kick(self, args): users_to_kick = [user for user in USERS if args[0] == user.ip_hash] if users_to_kick: await self.broadcast("KICKED", users_to_kick) for user in users_to_kick: await user.ws.send("YOU_KICKED") await user.ws.close(4000, "Has sido expulsado.") await self.ws.send("KICK_OK") print("Kicked " + args[0]) async def on_ban(self, args): users_to_ban = [user for user in USERS if args[0] == user.ip_hash] if users_to_ban: await self.broadcast("BANNED", users_to_ban) utils.add_ban(users_to_ban[0].ip, settings.bans_file) for user in users_to_ban: await user.ws.send("YOU_BANNED") await user.ws.close(4001, "Has sido baneado.") await self.ws.send("BAN_OK") print("Banned " + ip_to_ban) async def on_settitle(self, args): update_state('title', args[0]) await self.broadcast("TITLE:{}".format(args[0]), [self]) async def on_setsource(self, args): t = timestamp() update_state('source', args[0]) update_state('source_time', t) await self.broadcast("SOURCE:{}:{}".format(args[0], t), [self]) async def handle_message(self, msg_in): if len(msg_in) > settings.max_cmd_length: await 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: await cmd_to_run(args[1:]) else: await self.ws.send("FORBIDDEN") else: await cmd_to_run(args[1:]) else: await self.ws.send("INVALID") else: await self.ws.send("WHAT") async def broadcast(msg): print("broadcasting " + msg) if USERS: _ = [await user.ws.send(msg) for user in USERS] async def notify_count(): total_users = len(USERS) await broadcast("COUNT:{}".format(total_users)) async def register(websocket): USERS.add(websocket) await notify_count() async def unregister(websocket): USERS.remove(websocket) await notify_count() async def baitv_daemon(websocket, path): print("New client ({})".format(websocket.remote_address)) try: real_ip = websocket.request_headers['X-Real-IP'] except: real_ip = websocket.remote_address[0] this_client = Client(websocket, real_ip) await websocket.send("WELCOME:{}".format(VERSION)) #print(list(websocket.request_headers.raw_items())) if utils.is_banned(real_ip, settings.bans_file): await websocket.send("YOU_BANNED") await websocket.close(4002, "Estás baneado.") return if real_ip != "127.0.0.1": await register(this_client) try: while True: msg_in = await websocket.recv() print("< {}".format(msg_in)) await 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": await 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()