diff --git a/README.md b/README.md index 62a29ab..0ad8284 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,7 @@ https://photonglass.dev - **Multi Device Support**: Connect to multiple devices from one single interface. - **Easy Deployment**: Extremely easy to deploy and scale with multiple devices. - **Webhook Logging**: Log queries to a webhook channel (optional). +- **Rate Limiting**: Reduce service abuse by rate limiting users, 100 per day and 10 per minute. ## Setup (Docker) 1. Clone repository diff --git a/app/__init__.py b/app/__init__.py index 77aa092..5140f18 100644 --- a/app/__init__.py +++ b/app/__init__.py @@ -1,14 +1,25 @@ +import logging from flask import Flask -import logging, os +from flask_limiter import Limiter + +from app.functions.utils import get_client_ip logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) +limiter = Limiter( + get_client_ip, + default_limits=["100 per day", "10 per minute"], + storage_uri="memory://" +) + def create_app(): app = Flask(__name__, instance_path="/instance") + limiter.init_app(app) + from app.views import main app.register_blueprint(main.bp) app.add_url_rule('/', endpoint='index') - return app \ No newline at end of file + return app diff --git a/app/functions/utils.py b/app/functions/utils.py index 841bf6b..6dd1930 100644 --- a/app/functions/utils.py +++ b/app/functions/utils.py @@ -1,6 +1,7 @@ import os, yaml, logging, requests from functools import wraps from netmiko import ConnectHandler, NetmikoTimeoutException, NetmikoAuthenticationException +from flask import request logger = logging.getLogger(__name__) @@ -66,7 +67,7 @@ def send_webhook(webhook_url, text_data): # Get client IP address @exception_handler -def get_client_ip(request): +def get_client_ip(): if not request.headers.getlist("X-Forwarded-For"): return request.remote_addr return request.headers.getlist("X-Forwarded-For")[0] diff --git a/app/requirements.txt b/app/requirements.txt index e4d4d67..7cb26c1 100644 --- a/app/requirements.txt +++ b/app/requirements.txt @@ -3,4 +3,5 @@ pyyaml==6.0.2 netmiko==4.5.0 werkzeug==3.1.3 gunicorn==23.0.0 -requests==2.32.3 \ No newline at end of file +requests==2.32.3 +flask-limiter==3.9.2 \ No newline at end of file diff --git a/app/views/main.py b/app/views/main.py index 8cd123e..ef963a7 100644 --- a/app/views/main.py +++ b/app/views/main.py @@ -1,14 +1,15 @@ -from flask import ( - Blueprint, request, render_template -) import logging +from flask import Blueprint, request, render_template + from app.functions.utils import exception_handler, load_yaml, send_webhook, get_client_ip, execute_command logger = logging.getLogger(__name__) bp = Blueprint('main', __name__) + +# Route to render the main page @bp.route('/') @exception_handler def index():