import asyncio import json import logging from websockets.asyncio.server import serve from websockets.asyncio.server import broadcast import re from passlib.hash import scrypt import db import uuid import secrets import time version = "Helium-0.0.0a" addr = "localhost" port = 3636 logging.basicConfig(level=logging.INFO) error_contexts = { "malformedJson": "The JSON data sent to the server could not be parsed.", "lengthInvalid": "A value in the JSON data is longer or shorter than expected.", "invalidUsername": "Username is invalid. It may contain characters that are not permitted in usernames.", "invalidInvite": "The invite code you are trying to use is invalid or has expired.", "usernameTaken": "This username has been taken.", "notExists": "The requested value does not exist.", "lockdown": "Maintenance is in progress.", "authed": "You are already authenticated." } ulist = {} user_clients = {} clients = [] invite_codes = ["STANLEYYELNATSAB"] locked = False class util: def error(code, listener, data=None): if code in error_contexts: context = error_contexts[code] else: context = "" response = { "error": True, "code": code, "form": "helium-util", "context": context, "listener": listener } if data: response.update(data) return json.dumps(response) def field_check(expects, gets): for i in expects: if i not in gets: return "malformedJson" if type(gets[i]) in [str, dict, list]: if len(gets[i]) not in expects[i]: return "lengthInvalid" return True def greeting(): return json.dumps({ "command": "greet", "version": version, "ulist": ulist, "messages": [], "locked": locked }) def ulist(): broadcast(clients, json.dumps({ "command": "ulist", "ulist": ulist })) def author_data(username): data = db.acc.get(username) del data["secure"] del data["profile"] return data def authorize(username, conn_id, client=""): ulist[username] = client user_clients[conn_id] = {"username": username, "client": client} data = db.acc.get(username) del data["secure"] return data async def handler(websocket): clients.append(websocket) await websocket.send(util.greeting()) async for message in websocket: try: r = json.loads(message) except: await websocket.send(util.error("malformedJson", None)) continue if "listener" not in r: r["listener"] = None listener = r["listener"] if "command" not in r: await websocket.send(util.error("malformedJson", listener)) continue if r["command"] == "register": fc = util.field_check({"username": range(1,21), "password": range(8,256), "invite_code": range(16,17)}, r) if fc != True: await websocket.send(util.error(fc, listener)) continue if str(websocket.id) in user_clients: await websocket.send(util.error("authed", listener)) continue if locked: await websocket.send(util.error("lockdown", listener)) continue r["username"] = r["username"].lower() r["invite_code"] = r["invite_code"].upper() if not re.fullmatch("[a-z0-9-_]{1,20}", r["username"]): await websocket.send(util.error("invalidUsername", listener)) continue if r["invite_code"] not in invite_codes: await websocket.send(util.error("invalidInvite", listener)) continue if db.acc.get(r["username"]) != "notExists": await websocket.send(util.error("usernameTaken", listener)) continue data = { "_id": str(uuid.uuid4()), "username": r["username"], "display_name": r["username"], "created": round(time.time()), "avatar": None, "bot": False, "verified": False, "banned_until": 0, "permissions": [], "profile": { "bio": "", "lastfm": "", "banner": None, "links": {} }, "secure": { "password": scrypt.hash(r["password"]), "token": secrets.token_urlsafe(64), "ban_reason": "", "invite_code": r["invite_code"], "support_code": secrets.token_hex(16) } } result = db.acc.add(data) if result != True: await websocket.send(util.error(result, listener)) continue invite_codes.remove(r["invite_code"]) await websocket.send(json.dumps({"error": False, "token": data["secure"]["token"], "listener": listener})) elif r["command"] == "login_pswd": fc = util.field_check({"username": range(1,21), "password": range(8,256)}, r) if fc != True: await websocket.send(util.error(fc, listener)) continue if str(websocket.id) in user_clients: await websocket.send(util.error("authed", listener)) continue r["username"] = r["username"].lower() if locked: perms = db.acc.get_perms(r["username"]) if type(perms) != list: await websocket.send(util.error("lockdown", listener)) continue if "LOCK" not in perms: await websocket.send(util.error("lockdown", listener)) continue valid = db.acc.verify_pswd(r["username"], r["password"]) if type(valid) == dict: userdata = util.authorize(r["username"], str(websocket.id)) await websocket.send(json.dumps({"error": False, "token": valid["token"], "user": userdata, "listener": listener})) util.ulist() continue elif valid == "banned": await websocket.send(util.error(valid, listener, db.acc.get_ban(r["username"]))) continue else: await websocket.send(util.error(valid, listener)) continue elif r["command"] == "login_token": fc = util.field_check({"token": range(32,128)}, r) if fc != True: await websocket.send(util.error(fc, listener)) continue if str(websocket.id) in user_clients: await websocket.send(util.error("authed", listener)) continue if locked: await websocket.send(util.error("lockdown", listener)) continue valid = db.acc.verify(r["token"]) if type(valid) == dict: if valid["banned"]: await websocket.send(util.error("banned", listener, db.acc.get_ban(valid["username"]))) continue else: userdata = util.authorize(valid["username"], str(websocket.id)) await websocket.send(json.dumps({"error": False, "user": userdata, "listener": listener})) util.ulist() continue else: await websocket.send(util.error(valid, listener)) continue elif r["command"] == "ping": pass else: await websocket.send(util.error("malformedJson", listener)) if websocket in clients: clients.append(websocket) async def main(): async with serve(handler, addr, port) as server: await server.serve_forever() asyncio.run(main())