Add rate limiting support to API keys

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):

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):
         """

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;

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"))

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
             </a>
-                    <a hx-get="{{ url_for('settings.account') }}" hx-target="#container" hx-swap="innerHTML"
-                hx-push-url="true"
+            <a hx-get="{{ url_for('settings.account') }}" hx-target="#container" hx-swap="innerHTML" hx-push-url="true"
                 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">
                 Account
             </a>
@@ -44,24 +43,33 @@
         <table class="min-w-full divide-y divide-gray-200 dark:divide-gray-700">
             <thead class="bg-gray-50 dark:bg-gray-700">
                 <tr>
-                    <th
+                    <th scope="col"
                         class="px-6 py-3 text-left text-xs font-medium text-gray-500 dark:text-gray-300 uppercase tracking-wider">
-                        Name</th>
-                    <th
+                        Name
+                    </th>
+                    <th scope="col"
                         class="px-6 py-3 text-left text-xs font-medium text-gray-500 dark:text-gray-300 uppercase tracking-wider">
-                        Key Prefix</th>
-                    <th
+                        Key Prefix
+                    </th>
+                    <th scope="col"
                         class="px-6 py-3 text-left text-xs font-medium text-gray-500 dark:text-gray-300 uppercase tracking-wider">
-                        Scopes</th>
-                    <th
+                        Scopes
+                    </th>
+                    <th scope="col"
                         class="px-6 py-3 text-left text-xs font-medium text-gray-500 dark:text-gray-300 uppercase tracking-wider">
-                        Created</th>
-                    <th
+                        Rate Limit
+                    </th>
+                    <th scope="col"
                         class="px-6 py-3 text-left text-xs font-medium text-gray-500 dark:text-gray-300 uppercase tracking-wider">
-                        Last Used</th>
-                    <th
-                        class="px-6 py-3 text-right text-xs font-medium text-gray-500 dark:text-gray-300 uppercase tracking-wider">
-                        Actions</th>
+                        Created
+                    </th>
+                    <th scope="col"
+                        class="px-6 py-3 text-left text-xs font-medium text-gray-500 dark:text-gray-300 uppercase tracking-wider">
+                        Last Used
+                    </th>
+                    <th scope="col" class="relative px-6 py-3">
+                        <span class="sr-only">Delete</span>
+                    </th>
                 </tr>
             </thead>
             <tbody class="bg-white dark:bg-gray-800 divide-y divide-gray-200 dark:divide-gray-700">
@@ -72,11 +80,26 @@
                     <td class="px-6 py-4 whitespace-nowrap text-sm text-gray-500 dark:text-gray-400">{{ key.key[:8]
                         }}...</td>
                     <td class="px-6 py-4 text-sm text-gray-500 dark:text-gray-400">
-                        {% for scope in key.scopes %}
-                        <span
-                            class="inline-flex items-center px-2.5 py-0.5 rounded-full text-xs font-medium bg-blue-100 text-blue-800 dark:bg-blue-900 dark:text-blue-200 mr-1">{{
-                            scope }}</span>
-                        {% endfor %}
+                        <div class="text-sm text-gray-900 dark:text-white">
+                            {% if key.scopes == ['*'] %}
+                            <span
+                                class="px-2 inline-flex text-xs leading-5 font-semibold rounded-full bg-green-100 text-green-800 dark:bg-green-900 dark:text-green-200">Full
+                                Access</span>
+                            {% else %}
+                            <span
+                                class="px-2 inline-flex text-xs leading-5 font-semibold rounded-full bg-blue-100 text-blue-800 dark:bg-blue-900 dark:text-blue-200">{{
+                                key.scopes|length }} Scopes</span>
+                            {% endif %}
+                        </div>
+                    </td>
+                    <td class="px-6 py-4 whitespace-nowrap">
+                        <div class="text-sm text-gray-900 dark:text-white">
+                            {% if key.rate_limit_count %}
+                            {{ key.rate_limit_count }} / {{ key.rate_limit_period }}
+                            {% else %}
+                            <span class="text-gray-500 dark:text-gray-400">Unlimited</span>
+                            {% endif %}
+                        </div>
                     </td>
                     <td class="px-6 py-4 whitespace-nowrap text-sm text-gray-500 dark:text-gray-400">{{
                         key.created_at.strftime('%Y-%m-%d') }}</td>
@@ -91,7 +114,7 @@
                 </tr>
                 {% else %}
                 <tr>
-                    <td colspan="6" class="px-6 py-4 text-center text-sm text-gray-500 dark:text-gray-400">No API keys
+                    <td colspan="7" class="px-6 py-4 text-center text-sm text-gray-500 dark:text-gray-400">No API keys
                         found. Generate one to get started.</td>
                 </tr>
                 {% endfor %}
@@ -133,6 +156,27 @@
                 </div>
                 <p class="mt-1 text-xs text-gray-500">Select "Full Access" or specific functions.</p>
             </div>
+
+            <div class="mb-4">
+                <label class="block text-sm font-medium text-gray-700 dark:text-gray-300 mb-2">Rate Limit
+                    (Optional)</label>
+                <div class="flex space-x-4">
+                    <div class="w-1/2">
+                        <input type="number" name="rate_limit_count" placeholder="Requests count (e.g. 100)" min="1"
+                            class="w-full px-3 py-2 border border-gray-300 dark:border-gray-600 rounded-md shadow-sm focus:ring-blue-500 focus:border-blue-500 dark:bg-gray-700 dark:text-white">
+                    </div>
+                    <div class="w-1/2">
+                        <select name="rate_limit_period"
+                            class="w-full px-3 py-2 border border-gray-300 dark:border-gray-600 rounded-md shadow-sm focus:ring-blue-500 focus:border-blue-500 dark:bg-gray-700 dark:text-white">
+                            <option value="">No Limit</option>
+                            <option value="minute">Per Minute</option>
+                            <option value="hour">Per Hour</option>
+                            <option value="day">Per Day</option>
+                        </select>
+                    </div>
+                </div>
+            </div>
+
             <div class="flex justify-end space-x-3">
                 <button type="button" onclick="document.getElementById('create-key-modal').close()"
                     class="px-4 py-2 text-sm font-medium text-gray-700 bg-white border border-gray-300 rounded-md hover:bg-gray-50 dark:bg-gray-700 dark:text-gray-300 dark:border-gray-600 dark:hover:bg-gray-600">Cancel</button>
@@ -142,4 +186,4 @@
         </form>
     </div>
 </dialog>
-{% endblock %}
+{% endblock %}