path: root/main.py

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
from urllib.parse import urlparse

version = "Helium-0.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"
}

ulist = {}
client_data = {}

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": 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=""):
        ulist[username] = client
        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
                }))

    def forcekick(username):
        if username in ulist:
            del ulist[username]
        for i in client_data.copy():
            if client_data[i]["username"] == username:
                try:
                    client_data[i]["websocket"].close()
                except:
                    pass
                del client_data[i]

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 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)
                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)
                    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"] == "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"] == "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":
            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
            broadcast(clients, json.dumps({
                "command": "new_post",
                "data": data
            }))
        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))
    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())