src/utils/credential_manager.py

"""
Credential manager for secure API key handling.

Manages credentials for real API calls with validation and security.
"""

from typing import Optional, Dict
import os
from pathlib import Path
import json


class CredentialManager:
    """Manages API credentials securely."""
    
    def __init__(self, config_path: Optional[str] = None):
        """
        Initialize credential manager.
        
        Args:
            config_path: Optional path to credentials file
        """
        self.config_path = config_path or os.path.expanduser("~/.token-estimator/credentials.json")
        self.credentials: Dict[str, Dict] = {}
        self._load_credentials()
    
    def _load_credentials(self):
        """Load credentials from file if it exists."""
        config_file = Path(self.config_path)
        if config_file.exists():
            try:
                with open(config_file, 'r') as f:
                    self.credentials = json.load(f)
            except Exception as e:
                print(f"Warning: Could not load credentials: {e}")
                self.credentials = {}
        else:
            # Try environment variables
            self._load_from_env()
    
    def _load_from_env(self):
        """Load credentials from environment variables."""
        env_mapping = {
            'github': 'GITHUB_TOKEN',
            'slack': 'SLACK_TOKEN',
            'linear': 'LINEAR_API_KEY',
            'notion': 'NOTION_TOKEN',
            'stripe': 'STRIPE_API_KEY',
            'sentry': 'SENTRY_AUTH_TOKEN',
        }
        
        for service, env_var in env_mapping.items():
            value = os.getenv(env_var)
            if value:
                self.credentials[service] = {'token': value}
    
    def save_credentials(self):
        """Save credentials to file."""
        config_file = Path(self.config_path)
        config_file.parent.mkdir(parents=True, exist_ok=True)
        
        with open(config_file, 'w') as f:
            json.dump(self.credentials, f, indent=2)
        
        # Set restrictive permissions
        os.chmod(config_file, 0o600)
    
    def set_credential(self, service: str, credential_type: str, value: str):
        """
        Set a credential for a service.
        
        Args:
            service: Service name (e.g., 'github')
            credential_type: Type of credential (e.g., 'token', 'api_key')
            value: The credential value
        """
        if service not in self.credentials:
            self.credentials[service] = {}
        
        self.credentials[service][credential_type] = value
    
    def get_credential(self, service: str, credential_type: str = 'token') -> Optional[str]:
        """
        Get a credential for a service.
        
        Args:
            service: Service name
            credential_type: Type of credential to retrieve
        
        Returns:
            The credential value or None if not found
        """
        if service in self.credentials:
            return self.credentials[service].get(credential_type)
        return None
    
    def get_all_credentials(self, service: str) -> Optional[Dict]:
        """
        Get all credentials for a service.
        
        Args:
            service: Service name
        
        Returns:
            Dictionary of all credentials for the service
        """
        return self.credentials.get(service)
    
    def has_credentials(self, service: str) -> bool:
        """
        Check if credentials exist for a service.
        
        Args:
            service: Service name
        
        Returns:
            True if credentials exist
        """
        return service in self.credentials and len(self.credentials[service]) > 0
    
    def remove_credential(self, service: str):
        """
        Remove all credentials for a service.
        
        Args:
            service: Service name
        """
        if service in self.credentials:
            del self.credentials[service]
    
    def list_services(self) -> list[str]:
        """
        List all services with stored credentials.
        
        Returns:
            List of service names
        """
        return list(self.credentials.keys())
    
    def validate_format(self, service: str, credentials: Dict) -> tuple[bool, Optional[str]]:
        """
        Validate credential format for a service.
        
        Args:
            service: Service name
            credentials: Credentials dictionary
        
        Returns:
            (is_valid, error_message)
        """
        if service == 'github':
            token = credentials.get('token', '')
            valid_prefixes = ['ghp_', 'gho_', 'ghs_', 'github_pat_']
            if not any(token.startswith(p) for p in valid_prefixes):
                return False, "GitHub token should start with ghp_, gho_, ghs_, or github_pat_"
        
        elif service == 'slack':
            token = credentials.get('token', '')
            if not (token.startswith('xoxb-') or token.startswith('xoxp-')):
                return False, "Slack token should start with xoxb- (bot) or xoxp- (user)"
        
        elif service == 'linear':
            api_key = credentials.get('api_key', '')
            if not api_key.startswith('lin_api_'):
                return False, "Linear API key should start with lin_api_"
        
        elif service == 'notion':
            token = credentials.get('token', '')
            if not token.startswith('secret_'):
                return False, "Notion integration token should start with secret_"
        
        elif service == 'stripe':
            api_key = credentials.get('api_key', '')
            valid_prefixes = ['sk_test_', 'sk_live_', 'rk_test_', 'rk_live_']
            if not any(api_key.startswith(p) for p in valid_prefixes):
                return False, "Stripe API key should start with sk_test_, sk_live_, rk_test_, or rk_live_"
        
        elif service == 'sentry':
            token = credentials.get('token', '')
            if len(token) != 64:
                return False, "Sentry auth token should be 64 characters"
        
        return True, None
    
    def get_masked_credentials(self, service: str) -> Optional[Dict]:
        """
        Get credentials with values masked for display.
        
        Args:
            service: Service name
        
        Returns:
            Dictionary with masked credential values
        """
        creds = self.get_all_credentials(service)
        if not creds:
            return None
        
        masked = {}
        for key, value in creds.items():
            if len(value) > 8:
                masked[key] = value[:4] + '****' + value[-4:]
            else:
                masked[key] = '****'
        
        return masked


# Global instance
_credential_manager: Optional[CredentialManager] = None


def get_credential_manager() -> CredentialManager:
    """Get the global credential manager instance."""
    global _credential_manager
    if _credential_manager is None:
        _credential_manager = CredentialManager()
    return _credential_manager