Update app.py

app.py

@@ -4,31 +4,54 @@ import json
 from huggingface_hub import HfApi, create_repo, upload_file
 import random
 import string
+from flask_httpauth import HTTPTokenAuth
 
 # Import documentation from assembler_docs.py
 from assembler_docs import DOCUMENTATION
 
 app = Flask(__name__)
 
-# Hugging Face API token (set in Space settings)
+# Set up HTTP Token Authentication
+auth = HTTPTokenAuth(scheme='Bearer')
+
+# Hugging Face API token (set in Space settings or provided in request)
 HF_TOKEN = os.getenv("HF_TOKEN")
 if not HF_TOKEN:
     raise ValueError("HF_TOKEN not set. Add it in Space settings.")
 
 hf_api = HfApi()
 
+# Authentication function to validate the token
+@auth.verify_token
+def verify_token(token):
+    if not token:
+        return False
+    try:
+        # Attempt to verify the token by checking if it can access the user's profile or Space-Share
+        user_info = hf_api.whoami(token=token)
+        # Optionally, check if the token has write access to Space-Share
+        # This is a basic check; you might want to add more specific permission checks
+        return bool(user_info and "Space-Share" in user_info.get('orgs', [])) or token == HF_TOKEN
+    except Exception:
+        return False
+
 def generate_space_name():
     """Generate a unique Space name."""
     random_suffix = ''.join(random.choices(string.ascii_lowercase + string.digits, k=6))
     return f"GeneratedSpace-{random_suffix}"
 
 @app.route('/create-space', methods=['POST'])
+@auth.login_required
 def create_hf_space():
     try:
+        # Use the authenticated token for Hugging Face operations
+        token = auth.current_user()
+        hf_api = HfApi(token=token)
+
         # Parse JSON input
         data = request.get_json()
         if not data:
-            return jsonify({"error": "No JSON data provided"}), 400
+            return jsonify({"error": "No JSON data provided"}), 401
 
         # Extract parameters
         space_type = data.get("space_type", "gradio")  # Default to gradio if not specified
@@ -36,12 +59,12 @@ def create_hf_space():
         params = data.get("parameters", {})  # Optional parameters
 
         if not files:
-            return jsonify({"error": "No files provided in JSON"}), 400
+            return jsonify({"error": "No files provided in JSON"}), 401
 
         # Validate space_type
         valid_space_types = ["gradio", "static", "docker", "streamlit"]
         if space_type not in valid_space_types:
-            return jsonify({"error": f"Invalid space_type. Must be one of {valid_space_types}"}), 400
+            return jsonify({"error": f"Invalid space_type. Must be one of {valid_space_types}"}), 401
 
         # Create a unique Space name and repo under Space-Share namespace
         space_name = generate_space_name()
@@ -53,7 +76,7 @@ def create_hf_space():
                 repo_id=space_name,
                 repo_type="space",
                 space_sdk=space_type,
-                token=HF_TOKEN,
+                token=token,
                 private=False,
                 exist_ok=True  # Allow creation even if the repo might exist (rare case)
             )
@@ -76,7 +99,7 @@ def create_hf_space():
                     path_in_repo=filename,
                     repo_id=full_repo_id,
                     repo_type="space",
-                    token=HF_TOKEN
+                    token=token
                 )
             except Exception as e:
                 os.remove(f"temp_{filename}")
@@ -100,7 +123,7 @@ def create_hf_space():
                     path_in_repo="requirements.txt",
                     repo_id=full_repo_id,
                     repo_type="space",
-                    token=HF_TOKEN
+                    token=token
                 )
             except Exception as e:
                 os.remove("temp_requirements.txt")
@@ -125,7 +148,7 @@ CMD ["python", "app.py"]
                     path_in_repo="Dockerfile",
                     repo_id=full_repo_id,
                     repo_type="space",
-                    token=HF_TOKEN
+                    token=token
                 )
             except Exception as e:
                 os.remove("temp_Dockerfile")
@@ -140,11 +163,12 @@ CMD ["python", "app.py"]
         }), 200
 
     except json.JSONDecodeError:
-        return jsonify({"error": "Invalid JSON format"}), 400
+        return jsonify({"error": "Invalid JSON format"}), 401
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
 @app.route('/docs', methods=['GET'])
+@auth.login_required
 def get_docs():
     """Return the API documentation as plain text."""
     return Response(DOCUMENTATION, mimetype='text/plain'), 200