path: root/baitv-daemon.py

#!/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()