From d04b7f2120d3f9e99cb309abc6301907cf64c69e Mon Sep 17 00:00:00 2001 From: Peter Stockings Date: Tue, 2 Dec 2025 17:02:17 +1100 Subject: [PATCH] Add rate limiting support to API keys --- app.py | 4 + db.py | 84 ++++++++++++++----- migrations/005_add_api_key_rate_limiting.sql | 6 ++ routes/settings.py | 14 +++- templates/dashboard/settings/api_keys.html | 88 +++++++++++++++----- 5 files changed, 152 insertions(+), 44 deletions(-) create mode 100644 migrations/005_add_api_key_rate_limiting.sql diff --git a/app.py b/app.py index dea2af3..e6a8942 100644 --- a/app.py +++ b/app.py @@ -171,6 +171,10 @@ async def execute_http_function(user_id, function): api_key = db.get_api_key(api_key_value) if api_key and api_key['user_id'] == user_id: + # Check rate limit + if not db.check_and_increment_api_key_usage(api_key['id']): + return jsonify({'error': 'Rate limit exceeded'}), 429 + # Check Scopes scopes = api_key['scopes'] if isinstance(scopes, str): diff --git a/db.py b/db.py index 3ac7e38..73604fc 100644 --- a/db.py +++ b/db.py @@ -182,47 +182,89 @@ ORDER BY invocation_time DESC""", [http_function_id]) def update_user_theme_preference(self, user_id, theme): self.execute( 'UPDATE users SET theme_preference=%s WHERE id=%s', [theme, user_id], commit=True) - def get_http_function_history(self, function_id): http_function_history = self.execute( 'SELECT version_id, http_function_id, script_content, version_number, updated_at FROM http_functions_versions WHERE http_function_id=%s ORDER BY version_number DESC', [function_id]) return http_function_history - def create_api_key(self, user_id, name, key, scopes): - self.execute( - 'INSERT INTO api_keys (user_id, name, key, scopes) VALUES (%s, %s, %s, %s)', - [user_id, name, key, json.dumps(scopes)], - commit=True - ) + def create_api_key(self, user_id, name, key, scopes, rate_limit_count=None, rate_limit_period=None): + new_key = self.execute( + 'INSERT INTO api_keys (user_id, name, key, scopes, rate_limit_count, rate_limit_period) VALUES (%s, %s, %s, %s, %s, %s) RETURNING id, user_id, name, key, scopes, created_at, last_used_at, rate_limit_count, rate_limit_period', + [user_id, name, key, json.dumps(scopes), rate_limit_count, rate_limit_period], commit=True, one=True) + return new_key def get_api_key(self, key): api_key = self.execute( - 'SELECT id, user_id, name, key, scopes, created_at, last_used_at FROM api_keys WHERE key=%s', - [key], - one=True - ) + 'SELECT id, user_id, name, key, scopes, created_at, last_used_at, rate_limit_count, rate_limit_period, usage_count, usage_reset_at FROM api_keys WHERE key=%s', [key], one=True) + if api_key and api_key.get('scopes'): + if isinstance(api_key['scopes'], str): + api_key['scopes'] = json.loads(api_key['scopes']) return api_key def delete_api_key(self, user_id, key_id): self.execute( - 'DELETE FROM api_keys WHERE user_id=%s AND id=%s', - [user_id, key_id], - commit=True - ) + 'DELETE FROM api_keys WHERE user_id=%s AND id=%s', [user_id, key_id], commit=True) def list_api_keys(self, user_id): api_keys = self.execute( - 'SELECT id, user_id, name, key, scopes, created_at, last_used_at FROM api_keys WHERE user_id=%s ORDER BY created_at DESC', - [user_id] - ) + 'SELECT id, user_id, name, key, scopes, created_at, last_used_at, rate_limit_count, rate_limit_period, usage_count, usage_reset_at FROM api_keys WHERE user_id=%s ORDER BY created_at DESC', [user_id]) return api_keys def update_api_key_last_used(self, key_id): self.execute( - 'UPDATE api_keys SET last_used_at=NOW() WHERE id=%s', - [key_id], - commit=True + 'UPDATE api_keys SET last_used_at=NOW() WHERE id=%s', [key_id], commit=True) + + def check_and_increment_api_key_usage(self, key_id): + """ + Check if API key has exceeded rate limit and increment usage. + Returns True if allowed, False if rate limited. + """ + # Get current key data + key_data = self.execute( + 'SELECT rate_limit_count, rate_limit_period, usage_count, usage_reset_at FROM api_keys WHERE id=%s', + [key_id], one=True ) + + if not key_data or not key_data['rate_limit_count']: + return True # No limit set + + limit = key_data['rate_limit_count'] + period = key_data['rate_limit_period'] + usage = key_data['usage_count'] + reset_at = key_data['usage_reset_at'] + + import datetime + now = datetime.datetime.now() + + # Check if we need to reset the counter + if not reset_at or now >= reset_at: + # Calculate new reset time + if period == 'minute': + new_reset = now + datetime.timedelta(minutes=1) + elif period == 'hour': + new_reset = now + datetime.timedelta(hours=1) + elif period == 'day': + new_reset = now + datetime.timedelta(days=1) + else: + new_reset = now + datetime.timedelta(minutes=1) # Default + + # Reset usage and set new reset time + self.execute( + 'UPDATE api_keys SET usage_count=1, usage_reset_at=%s, last_used_at=NOW() WHERE id=%s', + [new_reset, key_id], commit=True + ) + return True + + # Check limit + if usage >= limit: + return False + + # Increment usage + self.execute( + 'UPDATE api_keys SET usage_count=usage_count+1, last_used_at=NOW() WHERE id=%s', + [key_id], commit=True + ) + return True def export_user_data(self, user_id): """ diff --git a/migrations/005_add_api_key_rate_limiting.sql b/migrations/005_add_api_key_rate_limiting.sql new file mode 100644 index 0000000..ddac030 --- /dev/null +++ b/migrations/005_add_api_key_rate_limiting.sql @@ -0,0 +1,6 @@ +-- Add rate limiting columns to api_keys table +ALTER TABLE api_keys +ADD COLUMN IF NOT EXISTS rate_limit_count INTEGER, +ADD COLUMN IF NOT EXISTS rate_limit_period VARCHAR(20), -- 'minute', 'hour', 'day' +ADD COLUMN IF NOT EXISTS usage_count INTEGER DEFAULT 0, +ADD COLUMN IF NOT EXISTS usage_reset_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP; diff --git a/routes/settings.py b/routes/settings.py index 7723c70..d5e440e 100644 --- a/routes/settings.py +++ b/routes/settings.py @@ -87,6 +87,18 @@ def create_api_key(): name = request.form.get("name", "My API Key") scopes_list = request.form.getlist("scopes") + rate_limit_count = request.form.get("rate_limit_count") + rate_limit_period = request.form.get("rate_limit_period") + + if rate_limit_count: + try: + rate_limit_count = int(rate_limit_count) + except ValueError: + rate_limit_count = None + + if not rate_limit_period in ['minute', 'hour', 'day']: + rate_limit_period = None + if not scopes_list: scopes = ["*"] else: @@ -95,7 +107,7 @@ def create_api_key(): # Generate a secure random key key = f"sk_{secrets.token_urlsafe(24)}" - db.create_api_key(user_id, name, key, scopes) + db.create_api_key(user_id, name, key, scopes, rate_limit_count, rate_limit_period) flash(f"API Key created: {key} - Save it now, you won't see it again!", "success") return redirect(url_for("settings.api_keys")) diff --git a/templates/dashboard/settings/api_keys.html b/templates/dashboard/settings/api_keys.html index 6bd213c..5440fe7 100644 --- a/templates/dashboard/settings/api_keys.html +++ b/templates/dashboard/settings/api_keys.html @@ -23,8 +23,7 @@ class="border-b-2 border-transparent text-gray-500 hover:text-gray-700 hover:border-gray-300 dark:text-gray-400 dark:hover:text-gray-300 py-4 px-1 text-sm font-medium cursor-pointer"> Login History - Account @@ -44,24 +43,33 @@ - - - - - - + Created + + + @@ -72,11 +80,26 @@ + @@ -91,7 +114,7 @@ {% else %} - {% endfor %} @@ -133,6 +156,27 @@

Select "Full Access" or specific functions.

+ +
+ +
+
+ +
+
+ +
+
+
+
@@ -142,4 +186,4 @@
-{% endblock %} +{% endblock %} \ No newline at end of file
- Name + - Key Prefix + - Scopes + - Created + - Last Used - Actions + Last Used + + Delete +
{{ key.key[:8] }}... - {% for scope in key.scopes %} - {{ - scope }} - {% endfor %} +
+ {% if key.scopes == ['*'] %} + Full + Access + {% else %} + {{ + key.scopes|length }} Scopes + {% endif %} +
+ 		+
+ {% if key.rate_limit_count %} + {{ key.rate_limit_count }} / {{ key.rate_limit_period }} + {% else %} + Unlimited + {% endif %} +
{{ key.created_at.strftime('%Y-%m-%d') }}
No API keys + No API keys found. Generate one to get started.