Update app.py

app.py

@@ -24,23 +24,35 @@ PRICING_PLANS = {
     'pro': {
         'name': 'Pro Plan',
         'price': 'Coming Soon',
-        'rate_limit': None # Unlimited
+        'rate_limit': None  # Unlimited
     }
 }
 
-# Define the directory for API keys
+# Define the directories for API keys
 API_KEYS_DIRECTORY = os.path.join(os.getcwd(), 'static', 'data')
+FREE_API_KEYS_DIRECTORY = os.path.join(API_KEYS_DIRECTORY, 'free')
+PAID_API_KEYS_DIRECTORY = os.path.join(API_KEYS_DIRECTORY, 'paid')
+
+# Ensure directories exist
+os.makedirs(FREE_API_KEYS_DIRECTORY, exist_ok=True)
+os.makedirs(PAID_API_KEYS_DIRECTORY, exist_ok=True)
 
 # Function to generate a new API key and save it to a file
-def generate_api_key(username):
+def generate_api_key():
     current_date = datetime.datetime.now().strftime("%Y%m%d")
+    username = ''.join(random.choices(string.ascii_lowercase + string.digits, k=4))  # Generate a 4-character username
     api_key = 'HUAI' + username + current_date + ''.join(random.choices(string.ascii_uppercase + string.digits, k=5))
+    
+    # Determine the directory based on the plan
+    api_keys_directory = FREE_API_KEYS_DIRECTORY if request.form.get('plan') == 'free' else PAID_API_KEYS_DIRECTORY
+    
     # Save the API key to a file
-    with open(os.path.join(API_KEYS_DIRECTORY, username + '.txt'), 'w') as file:
+    with open(os.path.join(api_keys_directory, username + '.txt'), 'w') as file:
         file.write(api_key)
     return api_key
 
-# Middleware to require an API key for each request
+# Middleware to require an API key for each request and enforce rate limits for free plan
+# Middleware to require an API key for each request and enforce rate limits for free plan
 def require_api_key(view_function):
     @wraps(view_function)
     def decorated_function(*args, **kwargs):
@@ -51,21 +63,50 @@ def require_api_key(view_function):
         if not api_key:
             api_key = request.args.get('HUAI')
 
+        if not api_key:
+            abort(401)  # Unauthorized
+        
         if not validate_api_key(api_key):
-            abort(401) # Unauthorized
+            abort(401)  # Unauthorized
+        
+        # Check if it's a free plan and enforce rate limit
+        if request.form.get('plan') == 'free':
+            username = api_key[4:8]  # Extract username from API key
+            if not check_rate_limit(username):
+                abort(429)  # Too Many Requests
+        
         return view_function(*args, **kwargs)
     return decorated_function
 
+
 # Function to validate an API key by checking if it matches any file in the directory
 def validate_api_key(api_key):
-    for filename in os.listdir(API_KEYS_DIRECTORY):
-        filepath = os.path.join(API_KEYS_DIRECTORY, filename)
+    api_keys_directory = FREE_API_KEYS_DIRECTORY if api_key.startswith('HUAI') else PAID_API_KEYS_DIRECTORY
+    for filename in os.listdir(api_keys_directory):
+        filepath = os.path.join(api_keys_directory, filename)
         with open(filepath, 'r') as file:
             if file.read().strip() == api_key:
                 return True
     return False
 
+# Function to check rate limit for free plan
+
+def check_rate_limit(username):
+    today_date = datetime.datetime.now().strftime("%Y%m%d")
+    # Count the number of requests made by the user today
+    requests_count = sum(1 for _ in os.listdir(FREE_API_KEYS_DIRECTORY) if _.startswith(username + '_' + today_date))
+    return requests_count
+
 # Routes with API key requirement
+@app.route('/api/usage', methods=['GET'])
+@require_api_key
+def get_api_usage():
+    # Extract username from API key
+    api_key = request.headers.get('HUSI') or request.args.get('HUAI')
+    username = api_key[4:8]
+    # Get the usage count for the user
+    usage_count = check_rate_limit(username)
+    return jsonify({'username': username, 'usage_count': usage_count})
 
 @app.route('/api/search', methods=['GET'])
 @require_api_key
@@ -183,26 +224,14 @@ def pricing():
 @app.route('/generate_key', methods=['GET', 'POST'])
 def generate_key():
     if request.method == 'POST':
-        username = request.form.get('username')
-        plan = request.form.get('plan', 'free') # Default to 'free' if not provided
-        
-        if not username:
-            return jsonify({'error': 'Username not provided'}), 400
-        
-        if plan not in PRICING_PLANS:
-            return jsonify({'error': 'Invalid plan'}), 400
-        
-        # Check if the user already has an API key
-        for filename in os.listdir(API_KEYS_DIRECTORY):
-            if filename.startswith(username):
-                return jsonify({'error': 'API key already exists for this username'}), 400
-        
-        # Generate a new API key
-        api_key = generate_api_key(username)
+        # Generate a new API key with a random username (userid)
+        api_key = generate_api_key()
         
         return jsonify({'api_key': api_key}), 201
     else:
         # Render the form for GET requests
         return render_template('index.html', plans=PRICING_PLANS)
+
 if __name__ == '__main__':
-    app.run(debug=True)
+    app.run(debug=True)
+