1 files changed, 69 insertions, 1 deletions
diff --git a/main.py b/main.py
index 39413c2..9446b71 100644
--- a/main.py
+++ b/main.py
@@ -3,22 +3,44 @@ 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
 
 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."
+}
+
 ulist = {}
-ulist_connection_ids = {}
+user_clients = {}
 clients = []
 
+invite_codes = ["STANLEYYELNATSAB"]
+locked = False
+
 class util:
     def error(code, listener):
+        if code in error_contexts:
+            context = error_contexts[code]
+        else:
+            context = ""
         return json.dumps({
             "error": True,
             "code": code,
             "form": "helium-util",
+            "context": context,
             "listener": listener
         })
     
@@ -50,6 +72,52 @@ async def handler(websocket):
             if fc != True:
                 await websocket.send(util.error(fc, 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"],
+                "avatar": None,
+                "bot": False,
+                "verified": False,
+                "banned_until": 0,
+                "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
+            r["username"] = r["username"].lower()
+
         else:
             await websocket.send(util.error("malformedJson", listener))
     if websocket in clients: