Upload auth.py with huggingface_hub

auth.py

@@ -0,0 +1,118 @@
+import os
+import hashlib
+import secrets
+from typing import Dict, Optional
+from datetime import datetime, timedelta
+import logging
+
+logger = logging.getLogger(__name__)
+
+class APIKeyManager:
+    """Manages API key authentication and rate limiting"""
+    
+    def __init__(self):
+        self.api_keys = {
+            os.getenv("API_KEY_1", "your-secure-api-key-1"): {
+                "user": "user1",
+                "created": datetime.now(),
+                "last_used": None,
+                "request_count": 0
+            },
+            os.getenv("API_KEY_2", "your-secure-api-key-2"): {
+                "user": "user2", 
+                "created": datetime.now(),
+                "last_used": None,
+                "request_count": 0
+            }
+        }
+        self.rate_limits = {}  # {api_key: {minute: count}}
+        self.max_requests_per_minute = int(os.getenv("RATE_LIMIT", "10"))
+    
+    def validate_api_key(self, api_key: str) -> Optional[str]:
+        """Validate API key and return user info"""
+        if api_key in self.api_keys:
+            self.api_keys[api_key]["last_used"] = datetime.now()
+            self.api_keys[api_key]["request_count"] += 1
+            return self.api_keys[api_key]["user"]
+        return None
+    
+    def check_rate_limit(self, api_key: str) -> bool:
+        """Check if API key has exceeded rate limit"""
+        current_minute = datetime.now().strftime("%Y-%m-%d-%H-%M")
+        
+        if api_key not in self.rate_limits:
+            self.rate_limits[api_key] = {}
+        
+        # Clean old entries (keep only last 5 minutes)
+        cutoff_time = datetime.now() - timedelta(minutes=5)
+        keys_to_remove = []
+        for minute_key in self.rate_limits[api_key]:
+            try:
+                minute_time = datetime.strptime(minute_key, "%Y-%m-%d-%H-%M")
+                if minute_time < cutoff_time:
+                    keys_to_remove.append(minute_key)
+            except ValueError:
+                keys_to_remove.append(minute_key)
+        
+        for key in keys_to_remove:
+            del self.rate_limits[api_key][key]
+        
+        # Check current minute
+        current_count = self.rate_limits[api_key].get(current_minute, 0)
+        if current_count >= self.max_requests_per_minute:
+            return False
+        
+        # Increment counter
+        self.rate_limits[api_key][current_minute] = current_count + 1
+        return True
+    
+    def get_api_key_stats(self, api_key: str) -> Optional[Dict]:
+        """Get statistics for an API key"""
+        if api_key in self.api_keys:
+            stats = self.api_keys[api_key].copy()
+            current_minute = datetime.now().strftime("%Y-%m-%d-%H-%M")
+            stats["current_minute_requests"] = self.rate_limits.get(api_key, {}).get(current_minute, 0)
+            stats["rate_limit"] = self.max_requests_per_minute
+            return stats
+        return None
+    
+    def generate_new_api_key(self, user: str) -> str:
+        """Generate a new secure API key"""
+        api_key = secrets.token_urlsafe(32)
+        self.api_keys[api_key] = {
+            "user": user,
+            "created": datetime.now(),
+            "last_used": None,
+            "request_count": 0
+        }
+        return api_key
+    
+    def revoke_api_key(self, api_key: str) -> bool:
+        """Revoke an API key"""
+        if api_key in self.api_keys:
+            del self.api_keys[api_key]
+            if api_key in self.rate_limits:
+                del self.rate_limits[api_key]
+            return True
+        return False
+    
+    def list_api_keys(self) -> Dict:
+        """List all API keys with their stats (without revealing the keys)"""
+        result = {}
+        for api_key, info in self.api_keys.items():
+            # Hash the API key for identification without revealing it
+            key_hash = hashlib.sha256(api_key.encode()).hexdigest()[:8]
+            result[key_hash] = {
+                "user": info["user"],
+                "created": info["created"].isoformat(),
+                "last_used": info["last_used"].isoformat() if info["last_used"] else None,
+                "request_count": info["request_count"]
+            }
+        return result
+
+# Global instance
+api_key_manager = APIKeyManager()
+
+def get_api_key_manager() -> APIKeyManager:
+    """Get the global API key manager instance"""
+    return api_key_manager