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