From d4237dd65e60dd2cfd48c1f111efb160eaa682b7 Mon Sep 17 00:00:00 2001
From: nsde <neostyxde@gmail.com>
Date: Wed, 30 Aug 2023 20:55:31 +0200
Subject: [PATCH] Small improvements

---
 api/handler.py         | 42 +++++++++++++++++++++++-------------------
 api/helpers/network.py | 14 ++++++++++++--
 api/main.py            |  2 +-
 screen.sh              | 11 +++++++----
 4 files changed, 43 insertions(+), 26 deletions(-)

diff --git a/api/handler.py b/api/handler.py
index e370ba5..bcbe01b 100644
--- a/api/handler.py
+++ b/api/handler.py
@@ -41,6 +41,8 @@ async def handle(incoming_request: fastapi.Request):
         payload = await incoming_request.json()
     except json.decoder.JSONDecodeError:
         payload = {}
+    except UnicodeDecodeError:
+        payload = {}
 
     received_key = incoming_request.headers.get('Authorization')
 
@@ -83,29 +85,31 @@ async def handle(incoming_request: fastapi.Request):
     if user['credits'] < cost:
         return await errors.error(429, 'Not enough credits.', 'Wait or earn more credits. Learn more on our website or Discord server.')
 
-    payload_with_vars = json.dumps(payload)
 
-    replace_dict = {
-        'timestamp': str(int(time.time())),
-        'date': time.strftime('%Y-%m-%d'),
-        'time': time.strftime('%H:%M:%S'),
-        'datetime': time.strftime('%Y-%m-%d %H:%M:%S'),
-        'model': payload.get('model', 'unknown'),
-    }
+    if not 'DISABLE_VARS' in key_tags:
+        payload_with_vars = json.dumps(payload)
 
-    if 'ALLOW_INSECURE_VARS' in key_tags:
-        replace_dict.update({
-            'my.ip': ip_address,
-            'my.id': str(user['_id']),
-            'my.role': user.get('role', 'default'),
-            'my.credits': str(user['credits']),
-            'my.discord': user.get('auth', {}).get('discord', ''),
-        })
+        replace_dict = {
+            'timestamp': str(int(time.time())),
+            'date': time.strftime('%Y-%m-%d'),
+            'time': time.strftime('%H:%M:%S'),
+            'datetime': time.strftime('%Y-%m-%d %H:%M:%S'),
+            'model': payload.get('model', 'unknown'),
+        }
 
-    for key, value in replace_dict.items():
-        payload_with_vars = payload_with_vars.replace(f'[[{key}]]', value)
+        if 'ALLOW_INSECURE_VARS' in key_tags:
+            replace_dict.update({
+                'my.ip': ip_address,
+                'my.id': str(user['_id']),
+                'my.role': user.get('role', 'default'),
+                'my.credits': str(user['credits']),
+                'my.discord': user.get('auth', {}).get('discord', ''),
+            })
 
-    payload = json.loads(payload_with_vars)
+        for key, value in replace_dict.items():
+            payload_with_vars = payload_with_vars.replace(f'[[{key}]]', value)
+
+        payload = json.loads(payload_with_vars)
 
     policy_violation = False
     if '/moderations' not in path:
diff --git a/api/helpers/network.py b/api/helpers/network.py
index 50859db..e4b4849 100644
--- a/api/helpers/network.py
+++ b/api/helpers/network.py
@@ -1,4 +1,9 @@
-import asyncio
+import os
+import time
+
+from dotenv import load_dotenv
+
+load_dotenv()
 
 async def get_ip(request) -> str:
     """Get the IP address of the incoming request."""
@@ -17,7 +22,7 @@ async def get_ip(request) -> str:
 
     return detected_ip
 
-def get_ip_sync(request) -> str:
+def get_ratelimit_key(request) -> str:
     """Get the IP address of the incoming request."""
 
     xff = None
@@ -32,4 +37,9 @@ def get_ip_sync(request) -> str:
 
     detected_ip = next((i for i in possible_ips if i), None)
 
+    for whitelisted_ip in os.getenv('NO_RATELIMIT_IPS', '').split():
+        if whitelisted_ip in detected_ip:
+            custom_key = f'whitelisted-{time.time()}'
+            return custom_key
+
     return detected_ip
diff --git a/api/main.py b/api/main.py
index 349854c..5e9bf09 100644
--- a/api/main.py
+++ b/api/main.py
@@ -32,7 +32,7 @@ app.include_router(core.router)
 
 limiter = Limiter(
     swallow_errors=True,
-    key_func=network.get_ip_sync, default_limits=[
+    key_func=network.get_ratelimit_key, default_limits=[
     '2/second',
     '20/minute',
     '300/hour'
diff --git a/screen.sh b/screen.sh
index c5db6aa..91c13cb 100755
--- a/screen.sh
+++ b/screen.sh
@@ -3,17 +3,20 @@
 # Commit to the production branch
 # git commit -am "Auto-trigger - Production server started" && git push origin Production
 
+# Kill production server
+fuser -k 2333/tcp
+
+# Clear production directory
+rm -rf /home/nova-prod/*
+
 # Copy files to production
 cp -r * /home/nova-prod
 
-# Copy env file to production
+# Copy .prod.env file to production
 cp env/.prod.env /home/nova-prod/.env
 
 # Change directory
 cd /home/nova-prod
 
-# Kill the production server
-fuser -k 2333/tcp
-
 # Start screen
 screen -S nova-api python run prod && sleep 5