api_key_encryption.py

"""
Simple API Key Encryption Module for Glossarion
Encrypts only specific API key fields including multi-key support
"""

import os
import json
import base64
from cryptography.fernet import Fernet
from pathlib import Path


class APIKeyEncryption:
    """Simple encryption handler for API keys"""
    
    def __init__(self):
        self.key_file = Path('.glossarion_key')
        self.cipher = self._get_or_create_cipher()
        
        # Define which fields to encrypt
        self.api_key_fields = [
            'api_key',
            'replicate_api_key',
            # Add more field names here if needed
        ]
    
    def _get_or_create_cipher(self):
        """Get existing cipher or create new one"""
        if self.key_file.exists():
            try:
                key = self.key_file.read_bytes()
                return Fernet(key)
            except:
                pass
        
        # Generate new key
        key = Fernet.generate_key()
        self.key_file.write_bytes(key)
        
        # Hide file on Windows
        if os.name == 'nt':
            import ctypes
            ctypes.windll.kernel32.SetFileAttributesW(str(self.key_file), 2)
        else:
            # Restrict permissions on Unix
            os.chmod(self.key_file, 0o600)
        
        return Fernet(key)
    
    def encrypt_value(self, value):
        """Encrypt a single value"""
        try:
            encrypted = self.cipher.encrypt(value.encode())
            return f"ENC:{base64.b64encode(encrypted).decode()}"
        except:
            return value
    
    def decrypt_value(self, value):
        """Decrypt a single value"""
        if not isinstance(value, str) or not value.startswith('ENC:'):
            return value
        
        try:
            encrypted_data = base64.b64decode(value[4:])
            return self.cipher.decrypt(encrypted_data).decode()
        except:
            return value
    
    def encrypt_multi_keys(self, multi_keys):
        """Encrypt API keys in multi_api_keys array"""
        if not isinstance(multi_keys, list):
            return multi_keys
        
        encrypted_keys = []
        for key_entry in multi_keys:
            if isinstance(key_entry, dict):
                encrypted_entry = key_entry.copy()
                # Encrypt the api_key field in each entry
                if 'api_key' in encrypted_entry and encrypted_entry['api_key']:
                    value = encrypted_entry['api_key']
                    if isinstance(value, str) and not value.startswith('ENC:'):
                        encrypted_entry['api_key'] = self.encrypt_value(value)
                encrypted_keys.append(encrypted_entry)
            else:
                encrypted_keys.append(key_entry)
        
        return encrypted_keys
    
    def decrypt_multi_keys(self, multi_keys):
        """Decrypt API keys in multi_api_keys array"""
        if not isinstance(multi_keys, list):
            return multi_keys
        
        decrypted_keys = []
        for key_entry in multi_keys:
            if isinstance(key_entry, dict):
                decrypted_entry = key_entry.copy()
                # Decrypt the api_key field in each entry
                if 'api_key' in decrypted_entry and decrypted_entry['api_key']:
                    decrypted_entry['api_key'] = self.decrypt_value(decrypted_entry['api_key'])
                decrypted_keys.append(decrypted_entry)
            else:
                decrypted_keys.append(key_entry)
        
        return decrypted_keys
    
    def encrypt_config(self, config):
        """Encrypt specific API key fields including multi-key support"""
        encrypted = config.copy()
        
        # Encrypt regular API key fields
        for field in self.api_key_fields:
            if field in encrypted and encrypted[field]:
                value = encrypted[field]
                # Only encrypt if not already encrypted
                if isinstance(value, str) and not value.startswith('ENC:'):
                    encrypted[field] = self.encrypt_value(value)
        
        # Encrypt multi_api_keys if present
        if 'multi_api_keys' in encrypted:
            encrypted['multi_api_keys'] = self.encrypt_multi_keys(encrypted['multi_api_keys'])
            
        # Encrypt fallback_keys if present
        if 'fallback_keys' in encrypted:
            encrypted['fallback_keys'] = self.encrypt_multi_keys(encrypted['fallback_keys'])
        
        return encrypted
    
    def decrypt_config(self, config):
        """Decrypt specific API key fields including multi-key support"""
        decrypted = config.copy()
        
        # Decrypt regular API key fields
        for field in self.api_key_fields:
            if field in decrypted and decrypted[field]:
                decrypted[field] = self.decrypt_value(decrypted[field])
        
        # Decrypt multi_api_keys if present
        if 'multi_api_keys' in decrypted:
            decrypted['multi_api_keys'] = self.decrypt_multi_keys(decrypted['multi_api_keys'])
            
        # Decrypt fallback_keys if present  
        if 'fallback_keys' in decrypted:
            decrypted['fallback_keys'] = self.decrypt_multi_keys(decrypted['fallback_keys'])
        
        return decrypted

# Simple interface functions
_handler = None

def get_handler():
    global _handler
    if _handler is None:
        _handler = APIKeyEncryption()
    return _handler

def encrypt_config(config):
    """Encrypt API keys in config"""
    return get_handler().encrypt_config(config)

def decrypt_config(config):
    """Decrypt API keys in config"""
    return get_handler().decrypt_config(config)

def migrate_config_file(config_file='config.json'):
    """Migrate existing config to encrypted format"""
    try:
        # Read config
        with open(config_file, 'r', encoding='utf-8') as f:
            config = json.load(f)
        
        # Check if already encrypted
        handler = get_handler()
        needs_encryption = False
        
        # Check regular API key fields
        for field in handler.api_key_fields:
            if field in config and config[field]:
                if isinstance(config[field], str) and not config[field].startswith('ENC:'):
                    needs_encryption = True
                    break
        
        # Check multi_api_keys
        if 'multi_api_keys' in config and isinstance(config['multi_api_keys'], list):
            for key_entry in config['multi_api_keys']:
                if isinstance(key_entry, dict) and 'api_key' in key_entry:
                    if key_entry['api_key'] and not key_entry['api_key'].startswith('ENC:'):
                        needs_encryption = True
                        break

        # Check fallback_keys
        if 'fallback_keys' in config and isinstance(config['fallback_keys'], list):
            for key_entry in config['fallback_keys']:
                if isinstance(key_entry, dict) and 'api_key' in key_entry:
                    if key_entry['api_key'] and not key_entry['api_key'].startswith('ENC:'):
                        needs_encryption = True
                        break
        
        if not needs_encryption:
            print("Config already encrypted or no API keys found.")
            return True
        
        # Backup
        backup_file = f"{config_file}.backup"
        with open(backup_file, 'w', encoding='utf-8') as f:
            json.dump(config, f, ensure_ascii=False, indent=2)
        print(f"Created backup: {backup_file}")
        
        # Encrypt
        encrypted = encrypt_config(config)
        
        # Save
        with open(config_file, 'w', encoding='utf-8') as f:
            json.dump(encrypted, f, ensure_ascii=False, indent=2)
        
        print("✅ Successfully encrypted API keys!")
        
        # Show summary
        if 'multi_api_keys' in config:
            print(f"  - Encrypted {len(config['multi_api_keys'])} multi-key entries")
        
        if 'fallback_keys' in config:
            print(f"  - Encrypted {len(config['fallback_keys'])} fallback-key entries")
    
        return True
        
    except Exception as e:
        print(f"❌ Error: {e}")
        return False


if __name__ == "__main__":
    # Simple migration script
    import sys
    
    config_file = 'config.json'
    if len(sys.argv) > 1:
        config_file = sys.argv[1]
    
    if os.path.exists(config_file):
        print(f"Encrypting API keys in {config_file}...")
        migrate_config_file(config_file)
    else:
        print(f"Config file not found: {config_file}")