import asyncio import json import logging from websockets.exceptions import ConnectionClosed 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 from urllib.parse import urlparse version = "Helium-1.0.0a" attachment_whitelist = ["u.cubeupload.com", "files.catbox.moe", "litter.catbox.moe", "i.ibb.co", "cubeupload.com", "media.tenor.com", "tenor.com", "c.tenor.com", "meower.fraudulent.loan", "fraudulent.loan", "deer.fraudulent.loan"] 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.", "invalidFormat": "Value contains invalid characters, or is too long.", "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.", "unauthorized": "You must be authorized to perform this action.", "deprecated": "This command is no longer supported." } deprecated = { "set_display_name": "set_property", "set_avatar": "set_property", "set_bio": "set_property", "set_lastfm": "set_property", "clear_home": "" } ulist = {} client_data = {} clients = [] invite_codes = [] 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": db.posts.get_recent(), "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, websocket, client): # TODO: statuses if client: if not re.fullmatch("[a-zA-Z0-9-_. ]{1,50}", client): client = "" ulist[username] = {"client": client, "status": ""} client_data[conn_id] = {"username": username, "client": client, "websocket": websocket} data = db.acc.get(username) del data["secure"] return data def loggedout(username, conn_id, websocket): del client_data[conn_id] if username in ulist: rm_user = True for i in client_data.copy(): if client_data[i]["username"] == username: rm_user = False if rm_user: del ulist[username] broadcast(clients, json.dumps({ "command": "ulist", "ulist": ulist })) async def forcekick(username): if username in ulist: del ulist[username] broadcast(clients, json.dumps({ "command": "ulist", "ulist": ulist })) for i in client_data.copy(): if client_data[i]["username"] == username: try: await client_data[i]["websocket"].close() except: pass if i in client_data: del client_data[i] async def handler(websocket): # make variables global to not cause problems later down the line global invite_codes clients.append(websocket) await websocket.send(util.greeting()) try: 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 client_data: 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 client_data: 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), websocket, r.get("client")) 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 client_data: 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), websocket, r.get("client")) 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"] == "get_user": fc = util.field_check({"username": range(1,21)}, r) if fc != True: await websocket.send(util.error(fc, listener)) continue if str(websocket.id) not in client_data: await websocket.send(util.error("unauthorized", listener)) continue data = db.acc.get(r["username"]) if type(data) != dict: await websocket.send(util.error(data, listener)) continue del data["secure"] await websocket.send(json.dumps({"error": False, "user": data, "listener": listener})) continue elif r["command"] == "get_post": fc = util.field_check({"id": range(8,128)}, r) if fc != True: await websocket.send(util.error(fc, listener)) continue if str(websocket.id) not in client_data: await websocket.send(util.error("unauthorized", listener)) continue data = db.posts.get_by_id(r["id"], True) if type(data) != dict: await websocket.send(util.error(data, listener)) continue await websocket.send(json.dumps({"error": False, "post": data, "listener": listener})) continue elif r["command"] == "set_property": fc = util.field_check({"property": range(1,64), "value": range(0,2048)}, r) if fc != True: await websocket.send(util.error(fc, listener)) continue if str(websocket.id) not in client_data: await websocket.send(util.error("unauthorized", listener)) continue username = client_data[str(websocket.id)]["username"] if r["property"] in ["bio", "lastfm"]: result = db.acc.edit({f"profile.{r['property']}": r["value"]}, username) if result: await websocket.send(json.dumps({"error": False, "listener": listener})) continue else: await websocket.send(util.error("fail", listener)) continue elif r["property"] == "display_name": if not re.fullmatch("[a-zA-Z0-9-_,:🅱️. ]{1,20}", r["value"]): await websocket.send(util.error("invalidFormat", listener)) continue result = db.acc.edit({r["property"]: r["value"]}, username) if result: await websocket.send(json.dumps({"error": False, "listener": listener})) continue else: await websocket.send(util.error("fail", listener)) continue elif r["property"] == "avatar": if not urlparse(r["value"]).hostname in attachment_whitelist: await websocket.send(util.error("invalidFormat", listener)) continue result = db.acc.edit({r["property"]: r["value"]}, username) if result: await websocket.send(json.dumps({"error": False, "listener": listener})) continue else: await websocket.send(util.error("fail", listener)) continue else: await websocket.send(util.error("malformedJson", listener)) continue elif r["command"] == "gen_invite": if str(websocket.id) not in client_data: await websocket.send(util.error("unauthorized", listener)) continue username = client_data[str(websocket.id)]["username"] if "INVITE" not in db.acc.get_perms(username): await websocket.send(util.error("unauthorized", listener)) continue new_invite = secrets.token_hex(8).upper() invite_codes.append(new_invite) await websocket.send(json.dumps({"error": False, "invite_code": new_invite, "invite_codes": invite_codes, "listener": listener})) continue elif r["command"] == "reset_invites": if str(websocket.id) not in client_data: await websocket.send(util.error("unauthorized", listener)) continue username = client_data[str(websocket.id)]["username"] if "INVITE" not in db.acc.get_perms(username): await websocket.send(util.error("unauthorized", listener)) continue invite_codes = [] await websocket.send(json.dumps({"error": False, "listener": listener})) continue elif r["command"] == "force_kick": fc = util.field_check({"username": range(1,21)}, r) if fc != True: await websocket.send(util.error(fc, listener)) continue if str(websocket.id) not in client_data: await websocket.send(util.error("unauthorized", listener)) continue username = client_data[str(websocket.id)]["username"] if "KICK" not in db.acc.get_perms(username): await websocket.send(util.error("unauthorized", listener)) continue if "PROTECTED" in db.acc.get_perms(r["username"].lower()): await websocket.send(util.error("unauthorized", listener)) continue try: await util.forcekick(r["username"].lower()) except Exception as e: raise e await websocket.send(json.dumps({"error": False, "listener": listener})) continue elif r["command"] == "ban": fc = util.field_check({"username": range(1,21), "banned_until": range(0,0), "ban_reason": range(0,2001)}, r) if fc != True: await websocket.send(util.error(fc, listener)) continue if str(websocket.id) not in client_data: await websocket.send(util.error("unauthorized", listener)) continue username = client_data[str(websocket.id)]["username"] if "KICK" not in db.acc.get_perms(username): await websocket.send(util.error("unauthorized", listener)) continue r["username"] = r["username"].lower() if "PROTECTED" in db.acc.get_perms(r["username"]): await websocket.send(util.error("unauthorized", listener)) continue ac = db.acc.edit({"banned_until": r["banned_until"]}, r["username"]) if ac != True: server.send_message(client, util.error(ac, listener)) ac = db.acc.edit({"secure.ban_reason": r["ban_reason"]}, r["username"]) if ac != True: server.send_message(client, util.error(ac, listener)) if r["username"] in ulist: await util.forcekick(r["username"]) await websocket.send(json.dumps({"error": False, "listener": listener})) continue elif r["command"] == "get_inbox": if str(websocket.id) not in client_data: await websocket.send(util.error("unauthorized", listener)) continue data = db.inbox.get_recent() await websocket.send(json.dumps({"error": False, "inbox": data, "listener": listener})) elif r["command"] == "post_inbox": fc = util.field_check({"content": range(0,3001), "attachments": range(0,4)}, r) if fc != True: await websocket.send(util.error(fc, listener)) continue if str(websocket.id) not in client_data: await websocket.send(util.error("unauthorized", listener)) continue username = client_data[str(websocket.id)]["username"] if "INBOX" not in db.acc.get_perms(username): await websocket.send(util.error("unauthorized", listener)) continue attachments = [] for i in r["attachments"]: if urlparse(i).hostname in attachment_whitelist: attachments.append(i) if len(r["content"]) == 0 and len(r["attachments"]) == 0: await websocket.send(util.error("lengthInvalid", listener)) continue data = { "_id": str(uuid.uuid4()), "created": round(time.time()), "content": r["content"], "attachments": attachments } posted = db.inbox.add(data) if posted != True: await websocket.send(util.error(fc, listener)) continue await websocket.send(json.dumps({"error": False, "listener": listener})) continue elif r["command"] == "post": fc = util.field_check({"content": range(0,3001), "replies": range(0,4), "attachments": range(0,4)}, r) if fc != True: await websocket.send(util.error(fc, listener)) continue if str(websocket.id) not in client_data: await websocket.send(util.error("unauthorized", listener)) continue attachments = [] for i in r["attachments"]: if urlparse(i).hostname in attachment_whitelist: attachments.append(i) if len(r["content"]) == 0 and len(r["attachments"]) == 0: await websocket.send(util.error("lengthInvalid", listener)) continue username = client_data[str(websocket.id)]["username"] author = util.author_data(username) replies = [] for i in r["replies"]: post = db.posts.get_by_id(i) if type(post) == dict: replies.append(post) data = { "_id": str(uuid.uuid4()), "created": round(time.time()), "content": r["content"], "replies": replies, "attachments": attachments, "author": username } posted = db.posts.add(data) if posted != True: await websocket.send(util.error(fc, listener)) continue data["author"] = author incr = -1 for j in data["replies"]: incr += 1 reply_author = db.acc.get_author(j["author"]) if type(reply_author) != dict: reply_author = {} data["replies"][incr]["author"] = reply_author broadcast(clients, json.dumps({ "command": "new_post", "data": data })) await websocket.send(json.dumps({"error": False, "listener": listener})) continue elif r["command"] == "ping": pass elif r["command"] in deprecated: await websocket.send(util.error("deprecated", listener, {"replacement": deprecated[r["command"]]})) continue else: await websocket.send(util.error("malformedJson", listener)) except ConnectionClosed: pass except Exception as e: raise e finally: if str(websocket.id) in client_data: util.loggedout(client_data[str(websocket.id)]["username"], str(websocket.id), websocket) if websocket in clients: clients.remove(websocket) async def main(): async with serve(handler, addr, port) as server: await server.serve_forever() asyncio.run(main())