enclava/backend/app/models/api_key.py

"""
API Key model
"""
from datetime import datetime, timedelta
from typing import Optional, Dict, Any, List
from sqlalchemy import Column, Integer, String, DateTime, Boolean, Text, JSON, ForeignKey
from sqlalchemy.orm import relationship
from app.db.database import Base


class APIKey(Base):
    """API Key model for authentication and access control"""
    
    __tablename__ = "api_keys"
    
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, nullable=False)  # Human-readable name for the API key
    key_hash = Column(String, unique=True, index=True, nullable=False)  # Hashed API key
    key_prefix = Column(String, index=True, nullable=False)  # First 8 characters for identification
    
    # User relationship
    user_id = Column(Integer, ForeignKey("users.id"), nullable=False)
    user = relationship("User", back_populates="api_keys")
    
    # Related data relationships
    budgets = relationship("Budget", back_populates="api_key", cascade="all, delete-orphan")
    usage_tracking = relationship("UsageTracking", back_populates="api_key", cascade="all, delete-orphan")
    
    # Key status and permissions
    is_active = Column(Boolean, default=True)
    permissions = Column(JSON, default=dict)  # Specific permissions for this key
    scopes = Column(JSON, default=list)  # OAuth-like scopes
    
    # Usage limits
    rate_limit_per_minute = Column(Integer, default=60)  # Requests per minute
    rate_limit_per_hour = Column(Integer, default=3600)  # Requests per hour
    rate_limit_per_day = Column(Integer, default=86400)  # Requests per day
    
    # Allowed resources
    allowed_models = Column(JSON, default=list)  # List of allowed LLM models
    allowed_endpoints = Column(JSON, default=list)  # List of allowed API endpoints
    allowed_ips = Column(JSON, default=list)  # IP whitelist
    allowed_chatbots = Column(JSON, default=list)  # List of allowed chatbot IDs for chatbot-specific keys
    
    # Budget configuration
    is_unlimited = Column(Boolean, default=True)  # Unlimited budget flag
    budget_limit_cents = Column(Integer, nullable=True)  # Budget limit in cents
    budget_type = Column(String, nullable=True)  # "total" or "monthly"
    
    # Metadata
    description = Column(Text, nullable=True)
    tags = Column(JSON, default=list)  # For organizing keys
    
    # Timestamps
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    last_used_at = Column(DateTime, nullable=True)
    expires_at = Column(DateTime, nullable=True)  # Optional expiration
    
    # Usage tracking
    total_requests = Column(Integer, default=0)
    total_tokens = Column(Integer, default=0)
    total_cost = Column(Integer, default=0)  # In cents
    
    # Relationships
    usage_tracking = relationship("UsageTracking", back_populates="api_key", cascade="all, delete-orphan")
    budgets = relationship("Budget", back_populates="api_key", cascade="all, delete-orphan")
    plugin_audit_logs = relationship("PluginAuditLog", back_populates="api_key")
    
    def __repr__(self):
        return f"<APIKey(id={self.id}, name='{self.name}', user_id={self.user_id})>"
    
    def to_dict(self, include_sensitive: bool = False):
        """Convert API key to dictionary for API responses"""
        data = {
            "id": self.id,
            "name": self.name,
            "key_prefix": self.key_prefix,
            "user_id": self.user_id,
            "is_active": self.is_active,
            "permissions": self.permissions,
            "scopes": self.scopes,
            "rate_limit_per_minute": self.rate_limit_per_minute,
            "rate_limit_per_hour": self.rate_limit_per_hour,
            "rate_limit_per_day": self.rate_limit_per_day,
            "allowed_models": self.allowed_models,
            "allowed_endpoints": self.allowed_endpoints,
            "allowed_ips": self.allowed_ips,
            "allowed_chatbots": self.allowed_chatbots,
            "description": self.description,
            "tags": self.tags,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
            "last_used_at": self.last_used_at.isoformat() if self.last_used_at else None,
            "expires_at": self.expires_at.isoformat() if self.expires_at else None,
            "total_requests": self.total_requests,
            "total_tokens": self.total_tokens,
            "total_cost_cents": self.total_cost,
            "is_unlimited": self.is_unlimited,
            "budget_limit": self.budget_limit_cents,  # Map to budget_limit for API response
            "budget_type": self.budget_type
        }
        
        if include_sensitive:
            data["key_hash"] = self.key_hash
        
        return data
    
    def is_expired(self) -> bool:
        """Check if the API key has expired"""
        if self.expires_at is None:
            return False
        return datetime.utcnow() > self.expires_at
    
    def is_valid(self) -> bool:
        """Check if the API key is valid and active"""
        return self.is_active and not self.is_expired()
    
    def has_permission(self, permission: str) -> bool:
        """Check if the API key has a specific permission"""
        return permission in self.permissions
    
    def has_scope(self, scope: str) -> bool:
        """Check if the API key has a specific scope"""
        return scope in self.scopes
    
    def can_access_model(self, model_name: str) -> bool:
        """Check if the API key can access a specific model"""
        if not self.allowed_models:  # Empty list means all models allowed
            return True
        return model_name in self.allowed_models
    
    def can_access_endpoint(self, endpoint: str) -> bool:
        """Check if the API key can access a specific endpoint"""
        if not self.allowed_endpoints:  # Empty list means all endpoints allowed
            return True
        return endpoint in self.allowed_endpoints
    
    def can_access_from_ip(self, ip_address: str) -> bool:
        """Check if the API key can be used from a specific IP"""
        if not self.allowed_ips:  # Empty list means all IPs allowed
            return True
        return ip_address in self.allowed_ips
    
    def can_access_chatbot(self, chatbot_id: str) -> bool:
        """Check if the API key can access a specific chatbot"""
        if not self.allowed_chatbots:  # Empty list means all chatbots allowed
            return True
        return chatbot_id in self.allowed_chatbots
    
    def update_usage(self, tokens_used: int = 0, cost_cents: int = 0):
        """Update usage statistics"""
        self.total_requests += 1
        self.total_tokens += tokens_used
        self.total_cost += cost_cents
        self.last_used_at = datetime.utcnow()
    
    def set_expiration(self, days: int):
        """Set expiration date in days from now"""
        self.expires_at = datetime.utcnow() + timedelta(days=days)
    
    def extend_expiration(self, days: int):
        """Extend expiration date by specified days"""
        if self.expires_at is None:
            self.expires_at = datetime.utcnow() + timedelta(days=days)
        else:
            self.expires_at = self.expires_at + timedelta(days=days)
    
    def revoke(self):
        """Revoke the API key"""
        self.is_active = False
        self.updated_at = datetime.utcnow()
    
    def add_scope(self, scope: str):
        """Add a scope to the API key"""
        if scope not in self.scopes:
            self.scopes.append(scope)
    
    def remove_scope(self, scope: str):
        """Remove a scope from the API key"""
        if scope in self.scopes:
            self.scopes.remove(scope)
    
    def add_allowed_model(self, model_name: str):
        """Add an allowed model"""
        if model_name not in self.allowed_models:
            self.allowed_models.append(model_name)
    
    def remove_allowed_model(self, model_name: str):
        """Remove an allowed model"""
        if model_name in self.allowed_models:
            self.allowed_models.remove(model_name)
    
    def add_allowed_endpoint(self, endpoint: str):
        """Add an allowed endpoint"""
        if endpoint not in self.allowed_endpoints:
            self.allowed_endpoints.append(endpoint)
    
    def remove_allowed_endpoint(self, endpoint: str):
        """Remove an allowed endpoint"""
        if endpoint in self.allowed_endpoints:
            self.allowed_endpoints.remove(endpoint)
    
    def add_allowed_ip(self, ip_address: str):
        """Add an allowed IP address"""
        if ip_address not in self.allowed_ips:
            self.allowed_ips.append(ip_address)
    
    def remove_allowed_ip(self, ip_address: str):
        """Remove an allowed IP address"""
        if ip_address in self.allowed_ips:
            self.allowed_ips.remove(ip_address)
    
    def add_allowed_chatbot(self, chatbot_id: str):
        """Add an allowed chatbot"""
        if chatbot_id not in self.allowed_chatbots:
            self.allowed_chatbots.append(chatbot_id)
    
    def remove_allowed_chatbot(self, chatbot_id: str):
        """Remove an allowed chatbot"""
        if chatbot_id in self.allowed_chatbots:
            self.allowed_chatbots.remove(chatbot_id)
    
    @classmethod
    def create_default_key(cls, user_id: int, name: str, key_hash: str, key_prefix: str) -> "APIKey":
        """Create a default API key with standard permissions"""
        return cls(
            name=name,
            key_hash=key_hash,
            key_prefix=key_prefix,
            user_id=user_id,
            is_active=True,
            permissions={
                "read": True,
                "write": True,
                "chat": True,
                "embeddings": True
            },
            scopes=[
                "chat.completions",
                "embeddings.create",
                "models.list"
            ],
            rate_limit_per_minute=60,
            rate_limit_per_hour=3600,
            rate_limit_per_day=86400,
            allowed_models=[],  # All models allowed by default
            allowed_endpoints=[],  # All endpoints allowed by default
            allowed_ips=[],  # All IPs allowed by default
            description="Default API key with standard permissions",
            tags=["default"]
        )
    
    @classmethod
    def create_restricted_key(cls, user_id: int, name: str, key_hash: str, key_prefix: str,
                            models: List[str], endpoints: List[str]) -> "APIKey":
        """Create a restricted API key with limited permissions"""
        return cls(
            name=name,
            key_hash=key_hash,
            key_prefix=key_prefix,
            user_id=user_id,
            is_active=True,
            permissions={
                "read": True,
                "chat": True
            },
            scopes=[
                "chat.completions"
            ],
            rate_limit_per_minute=30,
            rate_limit_per_hour=1800,
            rate_limit_per_day=43200,
            allowed_models=models,
            allowed_endpoints=endpoints,
            allowed_ips=[],
            description="Restricted API key with limited permissions",
            tags=["restricted"]
        )
    
    @classmethod
    def create_chatbot_key(cls, user_id: int, name: str, key_hash: str, key_prefix: str,
                          chatbot_id: str, chatbot_name: str) -> "APIKey":
        """Create a chatbot-specific API key"""
        return cls(
            name=name,
            key_hash=key_hash,
            key_prefix=key_prefix,
            user_id=user_id,
            is_active=True,
            permissions={
                "chatbot": True
            },
            scopes=[
                "chatbot.chat"
            ],
            rate_limit_per_minute=100,
            rate_limit_per_hour=6000,
            rate_limit_per_day=144000,
            allowed_models=[],  # Will use chatbot's configured model
            allowed_endpoints=[
                f"/api/v1/chatbot/external/{chatbot_id}/chat"
            ],
            allowed_ips=[],
            allowed_chatbots=[chatbot_id],
            description=f"API key for chatbot: {chatbot_name}",
            tags=["chatbot", f"chatbot-{chatbot_id}"]
        )