diff --git a/docker-compose.yml b/docker-compose.yml index f313714..097e236 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -163,6 +163,8 @@ services: - USERS_CONFIG_PATH=/app/configs/users.json - STORAGE_BACKEND=local # Options: local, s3, minio - STORAGE_BASE_PATH=/app/data + - RATE_LIMIT_STORAGE_URI=memory:// + - PROXY_COUNT=0 # Set to 1 when running behind a reverse proxy (nginx) volumes: - webserver_data_staged:/app/data/staged - webserver_data_archived:/app/data/archived diff --git a/webserver/main.py b/webserver/main.py index aede7ab..1dee237 100644 --- a/webserver/main.py +++ b/webserver/main.py @@ -26,6 +26,9 @@ login_required, current_user, ) +from flask_limiter import Limiter +from flask_limiter.util import get_remote_address +from werkzeug.middleware.proxy_fix import ProxyFix from werkzeug.security import check_password_hash from werkzeug.utils import secure_filename from datetime import datetime, timedelta, timezone @@ -37,6 +40,11 @@ app.secret_key = os.getenv("SECRET_KEY", os.urandom(32)) app.config["MAX_CONTENT_LENGTH"] = 500 * 1024 * 1024 # WSGI-level enforcement +# Trust X-Forwarded-For from this many upstream proxies (set to 1 when behind nginx) +_proxy_count = int(os.getenv("PROXY_COUNT", "0")) +if _proxy_count > 0: + app.wsgi_app = ProxyFix(app.wsgi_app, x_for=_proxy_count) + # ── User store (loaded from file at startup) ────────────────────────────────── _users_config_path = os.getenv("USERS_CONFIG_PATH", "/app/configs/users.json") try: @@ -50,6 +58,17 @@ login_manager.login_view = "login" login_manager.login_message = "Please log in to access this page." +# ── Rate Limiting ─────────────────────────────────────────────────────────── +limiter_storage_uri = os.getenv("RATE_LIMIT_STORAGE_URI", "memory://") +limiter = Limiter( + key_func=get_remote_address, + app=app, + default_limits=[], + storage_uri=limiter_storage_uri, + strategy="fixed-window", +) + + class User(UserMixin): def __init__(self, user_id: str): @@ -142,6 +161,10 @@ def user_db_scope(user_id: str): @app.route("/login", methods=["GET", "POST"]) +@limiter.limit( + "5 per minute", + error_message="Too many login attempts. Please try again later.", +) def login(): if current_user.is_authenticated: return redirect(url_for("overview")) @@ -882,6 +905,10 @@ def get_file_size_mb(filepath): @app.route("/api/upload", methods=["POST"]) +@limiter.limit( + "10 per minute", + error_message="Too many upload attempts. Please slow down.", +) @login_required def upload_file(): """Handle file upload via drag and drop""" @@ -946,6 +973,10 @@ def upload_file(): @app.route("/api/files", methods=["GET"]) +@limiter.limit( + "30 per minute", + error_message="Too many requests. Please slow down.", +) @login_required def get_files(): """Get list of files in data_incoming and data_staged_for_parsing directories""" @@ -1004,6 +1035,20 @@ def get_files(): # ==================== ERROR HANDLERS ==================== +@app.errorhandler(429) +def ratelimit_handler(e): + """Handle rate limit exceeded errors.""" + return ( + jsonify( + { + "error": "Rate limit exceeded", + "message": str(e.description), + } + ), + 429, + ) + + @app.errorhandler(404) def not_found(error): return render_template("404.html"), 404 diff --git a/webserver/requirements.txt b/webserver/requirements.txt index e9a5a70..f2394dc 100644 --- a/webserver/requirements.txt +++ b/webserver/requirements.txt @@ -1,5 +1,6 @@ Flask==2.3.3 Flask-Login==0.6.3 +Flask-Limiter==3.5.0 psycopg2-binary==2.9.7 folium==0.14.0 gunicorn==22.0.0 diff --git a/webserver/tests/test_grafana_proxy.py b/webserver/tests/test_grafana_proxy.py index 41afce8..902b424 100644 --- a/webserver/tests/test_grafana_proxy.py +++ b/webserver/tests/test_grafana_proxy.py @@ -23,10 +23,10 @@ def _make_grafana_response(status=200, content=b"ok", headers=None): @pytest.fixture def logged_in_client(monkeypatch): - """Test client with demo_openmrg actually logged in via the login route. + """Test client with login bypassed and current_user mocked. - Uses a real session so flask_login's current_user proxy resolves correctly - inside the grafana_proxy route handler. + Uses LOGIN_DISABLED to skip auth checks and mocks current_user + so flask_login's proxy resolves correctly inside grafana_proxy. """ monkeypatch.setitem( wm.USERS, @@ -36,10 +36,13 @@ def logged_in_client(monkeypatch): "display_name": "OpenMRG", }, ) + mock_user = Mock() + mock_user.id = "demo_openmrg" + mock_user.display_name = "OpenMRG" + monkeypatch.setattr(wm, "current_user", mock_user) + monkeypatch.setitem(wm.app.config, "LOGIN_DISABLED", True) wm.app.config["TESTING"] = True - client = wm.app.test_client() - client.post("/login", data={"username": "demo_openmrg", "password": "testpass"}) - return client + return wm.app.test_client() def test_grafana_proxy_injects_webauth_user_header(logged_in_client, monkeypatch):