suna/backend/credentials/credential_service.py

import os
import json
import uuid
import hashlib
import base64
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Dict, List, Any, Optional, Tuple

from cryptography.fernet import Fernet

from services.supabase import DBConnection
from utils.logger import logger


@dataclass(frozen=True)
class MCPCredential:
    credential_id: str
    account_id: str
    mcp_qualified_name: str
    display_name: str
    config: Dict[str, Any]
    is_active: bool
    last_used_at: Optional[datetime] = None
    created_at: Optional[datetime] = None
    updated_at: Optional[datetime] = None


@dataclass(frozen=True)
class MCPRequirement:
    qualified_name: str
    display_name: str
    enabled_tools: List[str] = field(default_factory=list)
    required_config: List[str] = field(default_factory=list)
    custom_type: Optional[str] = None


@dataclass
class CredentialRequest:
    account_id: str
    mcp_qualified_name: str
    display_name: str
    config: Dict[str, Any]


class CredentialNotFoundError(Exception):
    pass


class CredentialAccessDeniedError(Exception):
    pass


class EncryptionService:
    def __init__(self):
        self._encryption_key = self._get_or_create_encryption_key()
        self._cipher = Fernet(self._encryption_key)
    
    def _get_or_create_encryption_key(self) -> bytes:
        key_env = os.getenv("MCP_CREDENTIAL_ENCRYPTION_KEY")
        
        try:
            if isinstance(key_env, str):
                return key_env.encode('utf-8')
            else:
                return key_env
                
        except Exception as e:
            logger.error(f"Invalid encryption key: {e}")
            logger.warning("Generating new encryption key for this session")
            key = Fernet.generate_key()
            logger.info(f"Generated new encryption key. Set this in your environment:")
            logger.info(f"MCP_CREDENTIAL_ENCRYPTION_KEY={key.decode()}")
            return key
    
    def encrypt_config(self, config: Dict[str, Any]) -> Tuple[bytes, str]:
        config_json = json.dumps(config, sort_keys=True)
        config_bytes = config_json.encode('utf-8')
        
        config_hash = hashlib.sha256(config_bytes).hexdigest()
        encrypted_config = self._cipher.encrypt(config_bytes)
        
        return encrypted_config, config_hash
    
    def decrypt_config(self, encrypted_config: bytes, expected_hash: str) -> Dict[str, Any]:
        try:
            decrypted_bytes = self._cipher.decrypt(encrypted_config)
            
            actual_hash = hashlib.sha256(decrypted_bytes).hexdigest()
            if actual_hash != expected_hash:
                raise ValueError("Credential integrity check failed")
            
            config_json = decrypted_bytes.decode('utf-8')
            return json.loads(config_json)
            
        except Exception as e:
            logger.error(f"Failed to decrypt credential: {e}")
            raise ValueError("Failed to decrypt credential")


class CredentialService:
    def __init__(self, db_connection: DBConnection):
        self._db = db_connection
        self._encryption = EncryptionService()
    
    async def store_credential(
        self,
        account_id: str,
        mcp_qualified_name: str,
        display_name: str,
        config: Dict[str, Any]
    ) -> str:
        logger.info(f"Storing credential for {mcp_qualified_name}")
        
        credential_id = str(uuid.uuid4())
        encrypted_config, config_hash = self._encryption.encrypt_config(config)
        encoded_config = base64.b64encode(encrypted_config).decode('utf-8')
        
        client = await self._db.client
        
        existing = await client.table('user_mcp_credentials').select('credential_id')\
            .eq('account_id', account_id)\
            .eq('mcp_qualified_name', mcp_qualified_name)\
            .eq('is_active', True)\
            .execute()
        
        if existing.data:
            await client.table('user_mcp_credentials').update({
                'is_active': False,
                'updated_at': datetime.now(timezone.utc).isoformat()
            }).eq('credential_id', existing.data[0]['credential_id']).execute()
        
        result = await client.table('user_mcp_credentials').insert({
            'credential_id': credential_id,
            'account_id': account_id,
            'mcp_qualified_name': mcp_qualified_name,
            'display_name': display_name,
            'encrypted_config': encoded_config,
            'config_hash': config_hash,
            'is_active': True,
            'created_at': datetime.now(timezone.utc).isoformat(),
            'updated_at': datetime.now(timezone.utc).isoformat()
        }).execute()
        
        logger.info(f"Stored credential {credential_id} for {mcp_qualified_name}")
        return credential_id
    
    async def get_credential(
        self, 
        account_id: str, 
        mcp_qualified_name: str
    ) -> Optional[MCPCredential]:
        client = await self._db.client
        result = await client.table('user_mcp_credentials').select('*')\
            .eq('account_id', account_id)\
            .eq('mcp_qualified_name', mcp_qualified_name)\
            .eq('is_active', True)\
            .execute()
        
        if not result.data:
            return None
        
        return self._map_to_credential(result.data[0])
    
    async def get_user_credentials(self, account_id: str) -> List[MCPCredential]:
        client = await self._db.client
        result = await client.table('user_mcp_credentials').select('*')\
            .eq('account_id', account_id)\
            .eq('is_active', True)\
            .order('created_at', desc=True)\
            .execute()
        
        return [self._map_to_credential(data) for data in result.data]
    
    async def delete_credential(
        self, 
        account_id: str, 
        mcp_qualified_name: str
    ) -> bool:
        logger.info(f"Deleting credential for {mcp_qualified_name}")
        
        client = await self._db.client
        result = await client.table('user_mcp_credentials').update({
            'is_active': False,
            'updated_at': datetime.now(timezone.utc).isoformat()
        }).eq('account_id', account_id)\
          .eq('mcp_qualified_name', mcp_qualified_name)\
          .eq('is_active', True)\
          .execute()
        
        success = len(result.data) > 0
        if success:
            logger.info(f"Deleted credential for {mcp_qualified_name}")
        
        return success
    
    async def get_missing_credentials(
        self, 
        account_id: str, 
        requirements: List[MCPRequirement]
    ) -> List[MCPRequirement]:
        user_credentials = await self.get_user_credentials(account_id)
        credential_names = {cred.mcp_qualified_name for cred in user_credentials}
        
        missing = []
        for req in requirements:
            if req.qualified_name not in credential_names:
                missing.append(req)
        
        return missing
    
    async def build_credential_mappings(
        self, 
        account_id: str, 
        requirements: List[MCPRequirement]
    ) -> Dict[str, str]:
        mappings = {}
        
        for req in requirements:
            if req.custom_type:
                user_credentials = await self.get_user_credentials(account_id)
                custom_pattern = f"custom_{req.custom_type}_"
                
                for cred in user_credentials:
                    if (cred.mcp_qualified_name.startswith(custom_pattern) and 
                        req.display_name.lower().replace(' ', '_') in cred.mcp_qualified_name):
                        mappings[req.qualified_name] = cred.credential_id
                        break
            else:
                credential = await self.get_credential(account_id, req.qualified_name)
                if credential:
                    mappings[req.qualified_name] = credential.credential_id
        
        return mappings
    
    async def validate_access(self, credential: MCPCredential, account_id: str) -> None:
        if credential.account_id != account_id:
            raise CredentialAccessDeniedError("Access denied to credential")
    
    def _map_to_credential(self, data: Dict[str, Any]) -> MCPCredential:
        try:
            encrypted_config = base64.b64decode(data['encrypted_config'])
            config = self._encryption.decrypt_config(encrypted_config, data['config_hash'])
        except Exception as e:
            logger.error(f"Failed to decrypt credential {data['credential_id']}: {e}")
            config = {}
        
        return MCPCredential(
            credential_id=data['credential_id'],
            account_id=data['account_id'],
            mcp_qualified_name=data['mcp_qualified_name'],
            display_name=data['display_name'],
            config=config,
            is_active=data['is_active'],
            last_used_at=datetime.fromisoformat(data['last_used_at'].replace('Z', '+00:00')) if data.get('last_used_at') else None,
            created_at=datetime.fromisoformat(data['created_at'].replace('Z', '+00:00')) if data.get('created_at') else None,
            updated_at=datetime.fromisoformat(data['updated_at'].replace('Z', '+00:00')) if data.get('updated_at') else None
        )


def get_credential_service(db_connection: DBConnection) -> CredentialService:
    return CredentialService(db_connection)