socket_handler.py

#!/usr/bin/env python

#Tornado
import tornado.websocket
import json
import string
import sys

from random import randint
from time import time, sleep

idLength = 10

#Handles the websocket requests
class SocketHandler(tornado.websocket.WebSocketHandler):
    clients = []
    maxClients = 4 

    logMode = "on"

    """
    Player handling methods
    """

    def getID() -> str:
        characters = string.ascii_letters + string.digits
        returnVal = ""
        for _ in range(0, idLength):
            returnVal += characters[randint(0, len(characters)-1)]

        return returnVal

    def getUID() -> str:
        """Get a unique id, consisting of lower and upprcase english letters + numbers from 0 to 9. The length is determined by the 'idLength' variable"""
        newId = SocketHandler.getID()
        while newId in SocketHandler.clients:
            newId = SocketHandler.getID()

        return newId

    # Sets the game data by socket
    def setUserData(socket, data):
        match = list( filter(lambda a : a[1]["socket"] == socket, enumerate(SocketHandler.clients)) )
        if len(match) == 0: return
        else:
            index = match[0][0]
            SocketHandler.clients[index]["data"] = data

    # Gets the game data by socket
    def getUserData(socket):
        match = list( filter(lambda a : a["socket"] == socket, SocketHandler.clients) )
        if len(match) == 0: return None
        else: return match[0]["data"]

    # Gets the game data by id
    def getUserDataByID(id):
        match = list( filter(lambda a : a["data"]["id"] == id, SocketHandler.clients) )
        if len(match) == 0: return None
        else: return match[0]["data"]

    def removeUser(socket):
        return list( filter(lambda a : a["socket"] != socket, SocketHandler.clients) )

    """
    Server methods
    """

    # Closesd all terminals and disconnects all users from the server
    def fullStop():
        for client in SocketHandler.clients:
            if "data" in client and "term" in client["data"]:
                client["data"]["term"].stop()

            client["socket"].close(reason = f"Server was shut down!")

            SocketHandler.removeUser(client["socket"])


    """
    Socket handling methods
    """

    def check_origin(self, origin) -> bool:
        return True

    def open(self):
        if len(SocketHandler.clients) >= SocketHandler.maxClients:
            if SocketHandler.logMode != "off": #on / reduced
                print(f"\rMaximum number of clients reached! ({SocketHandler.maxClients})")

            #raise tornado.web.HTTPError(403, "Connection refused.")
            self.close(reason = "Maximum number of clients reached")
            return
    
        newId = SocketHandler.getUID()

        # Add user to clients list, send them their ID
        SocketHandler.clients.append({
            "socket": self,
            "data": {
                "id": newId,
            }
        })

    def on_message(self, message):
        global GLOBAL_PASSWORD

        msg = None

        try:
            msg = json.loads(message)
        except:
            return

        if not "type" in msg: return

        # Somebody executed a command or pressed a key
        if msg["type"] == "message":
            userData = SocketHandler.getUserDataByID(msg["id"])
            if userData == None: return

            msg["username"] = userData["username"]

            # Add timestamp & broadcast message to all clients
            msg["time"] = time()
            for c in SocketHandler.clients:
                c["socket"].write_message(msg)

            return

        if msg["type"] == "connect":
            # Validate input
            if (not "password" in msg) or msg["password"] == "":
                self.write_message({"type": "validate", "state": "fail", "reason": "Password cannot be empty!"})
                return
    
            # Validate input
            if (not "username" in msg) or msg["username"] == "":
                self.write_message({"type": "validate", "state": "fail", "reason": "Username cannot be empty!"})
                return

            # Validate input
            if msg["password"] != sys.argv[1]:
                self.write_message({"type": "validate", "state": "fail", "reason": "Incorrect password!"})
                return

            # Get user ID
            userData = SocketHandler.getUserData(self)

            if SocketHandler.logMode == "on":
                print(f"\rUser '{msg['username']}' connected")

            # Let the client know, it registered successfully, and send them their ID
            self.write_message({
                "type": "validate",
                "id": userData["id"],
                "state": "pass",
                "username": msg["username"],
            })

            SocketHandler.setUserData(self, {
                "id": userData["id"], # Keep the user id
                "login_time": time(),
                "username": msg["username"],
            })

            return

    def on_close(self):
        print(f"\rUser left")

        SocketHandler.clients = SocketHandler.removeUser(self)