1 files changed, 40 insertions, 2 deletions
diff --git a/main.py b/main.py
index e840229..d1e321c 100644
--- a/main.py
+++ b/main.py
@@ -60,6 +60,7 @@ client_data = {}
 ratelimits = {}
 
 clients = []
+ips_by_client = {}
 
 invite_codes = []
 locked = False
@@ -144,6 +145,8 @@ class util:
                 client = ""
         ulist[username] = {"client": client, "status": "", "bot": bot}
         client_data[conn_id] = {"username": username, "client": client, "websocket": websocket, "connected": time.time(), "bot": bot}
+        if ips_by_client[websocket]:
+            db.acc.add_ip(ips_by_client[websocket], username)
         data = db.acc.get(username)
         del data["secure"]
         return data
@@ -189,6 +192,10 @@ async def handler(websocket):
     global invite_codes
 
     clients.append(websocket)
+    try:
+        ips_by_client[websocket] = websocket.request.headers["CF-Connecting-IP"]
+    except:
+        ips_by_client[websocket] = None
     ratelimits[str(websocket.id)] = 0
     await websocket.send(util.greeting())
     try:
@@ -197,7 +204,15 @@ async def handler(websocket):
             print(time.time())
             print(time.time() > ratelimits[str(websocket.id)])
             if not ratelimits[str(websocket.id)] <= time.time():
-                await websocket.send(util.error("ratelimited", None))
+                lst = None
+                try:
+                    r = json.loads(message)
+                    if "listener" not in r:
+                        r["listener"] = None
+                    lst = r["listener"]
+                except:
+                    pass
+                await websocket.send(util.error("ratelimited", lst))
                 continue
             ratelimits[str(websocket.id)] = time.time() + 0.025
             try:
@@ -233,6 +248,9 @@ async def handler(websocket):
                 if db.acc.get(r["username"]) != "notExists":
                     await websocket.send(util.error("usernameTaken", listener))
                     continue
+                ips = []
+                if ips_by_client[websocket]:
+                    ips.append(ips_by_client[websocket])
                 data = {
                     "_id": str(uuid.uuid4()),
                     "username": r["username"],
@@ -253,7 +271,8 @@ async def handler(websocket):
                         "token": secrets.token_urlsafe(64),
                         "ban_reason": "",
                         "invite_code": r["invite_code"],
-                        "support_code": secrets.token_hex(16)
+                        "support_code": secrets.token_hex(16),
+                        "ips": ips
                     }
                 }
                 result = db.acc.add(data)
@@ -456,6 +475,25 @@ async def handler(websocket):
                     await util.forcekick(r["username"])
                 await websocket.send(json.dumps({"error": False, "listener": listener}))
                 continue
+            elif r["command"] == "get_ips":
+                fc = util.field_check({"username": {"range": range(1,21), "types": [str]}}, 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 "IP" not in db.acc.get_perms(username):
+                    await websocket.send(util.error("unauthorized", listener))
+                    continue
+                r["username"] = r["username"].lower()
+                ac = db.acc.get(r["username"])
+                if ac == "notExists":
+                    server.send_message(client, util.error(ac, listener))
+                ac = ac["secure"]["ips"]
+                await websocket.send(json.dumps({"error": False, "ips": ac, "listener": listener}))
+                continue
             elif r["command"] == "get_inbox":
                 if str(websocket.id) not in client_data:
                     await websocket.send(util.error("unauthorized", listener))