ratelimit

1 files changed, 14 insertions, 1 deletions

diff --git a/main.py b/main.py
index b1f709f..19035df 100644
--- a/main.py
+++ b/main.py
@@ -31,7 +31,8 @@ error_contexts = {
     "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": "This command is no longer supported.",

+    "ratelimited": "You are ratelimited."

 }

 

 deprecated = {

@@ -45,6 +46,8 @@ deprecated = {
 ulist = {}

 client_data = {}

 

+ratelimits = {}

+

 clients = []

 

 invite_codes = []

@@ -143,9 +146,17 @@ async def handler(websocket):
     global invite_codes

 

     clients.append(websocket)

+    ratelimits[str(websocket.id)] = 0

     await websocket.send(util.greeting())

     try:

         async for message in websocket:

+            print(ratelimits[str(websocket.id)])

+            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))

+                continue

+            ratelimits[str(websocket.id)] = time.time() + 0.35

             try:

                 r = json.loads(message)

             except:

@@ -556,6 +567,8 @@ async def handler(websocket):
     finally:

         if str(websocket.id) in client_data:

             util.loggedout(client_data[str(websocket.id)]["username"], str(websocket.id), websocket)

+        if str(websocket.id) in ratelimits:

+            del ratelimits[str(websocket.id)]

         if websocket in clients:

             clients.remove(websocket)