File size: 3,884 Bytes
ec8f374
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
"""
Secure Configuration - Encrypted API key storage using Fernet encryption

Provides secure storage and retrieval of API keys with encryption at rest.
"""

import os
import json
from pathlib import Path
from typing import Dict, Optional
from cryptography.fernet import Fernet


class SecureConfig:
    """Manages encrypted API key storage"""

    def __init__(self, secrets_dir: str = ".secrets"):
        self.secrets_dir = Path(secrets_dir)
        self.secrets_dir.mkdir(exist_ok=True)

        self.key_file = self.secrets_dir / "encryption.key"
        self.data_file = self.secrets_dir / "api_keys.enc"

        self.fernet = self._load_or_create_key()

    def _load_or_create_key(self) -> Fernet:
        """Load existing encryption key or create new one"""
        if self.key_file.exists():
            with open(self.key_file, "rb") as f:
                key = f.read()
        else:
            key = Fernet.generate_key()
            with open(self.key_file, "wb") as f:
                f.write(key)

        return Fernet(key)

    def save_keys(self, api_keys: Dict[str, str]):
        """Save API keys with encryption"""
        # Remove empty keys
        api_keys = {k: v for k, v in api_keys.items() if v and v.strip()}

        # Serialize to JSON
        json_data = json.dumps(api_keys)

        # Encrypt
        encrypted_data = self.fernet.encrypt(json_data.encode())

        # Save to file
        with open(self.data_file, "wb") as f:
            f.write(encrypted_data)

        # Also set environment variables for current session
        for key, value in api_keys.items():
            os.environ[key] = value

    def load_keys(self) -> Dict[str, str]:
        """Load and decrypt API keys"""
        if not self.data_file.exists():
            return {}

        try:
            # Read encrypted data
            with open(self.data_file, "rb") as f:
                encrypted_data = f.read()

            # Decrypt
            decrypted_data = self.fernet.decrypt(encrypted_data)

            # Deserialize
            api_keys = json.loads(decrypted_data.decode())

            # Set environment variables
            for key, value in api_keys.items():
                os.environ[key] = value

            return api_keys

        except Exception as e:
            print(f"Error loading keys: {e}")
            return {}

    def get_key(self, key_name: str) -> Optional[str]:
        """Get a specific API key"""
        # First try environment variable
        value = os.environ.get(key_name)
        if value:
            return value

        # Then try loading from encrypted file
        keys = self.load_keys()
        return keys.get(key_name)

    def get_masked_key(self, key_name: str) -> Optional[str]:
        """Get a masked version of an API key for display"""
        key = self.get_key(key_name)
        if not key:
            return None

        # Show first 8 and last 4 characters
        if len(key) > 20:
            return f"{key[:8]}...{key[-4:]}"
        else:
            return f"{key[:4]}...{key[-2:]}"

    def delete_keys(self):
        """Delete all stored keys"""
        if self.data_file.exists():
            self.data_file.unlink()

        # Also clear environment variables
        env_keys = [
            "HUGGINGFACE_TOKEN",
            "OPENAI_API_KEY",
            "ANTHROPIC_API_KEY",
            "WANDB_API_KEY",
            "ELEVEN_API_KEY",
            "RUNPOD_API_KEY"
        ]
        for key in env_keys:
            os.environ.pop(key, None)

    def has_key(self, key_name: str) -> bool:
        """Check if a key exists"""
        return self.get_key(key_name) is not None


# Global instance
_config = None


def get_config() -> SecureConfig:
    """Get the global secure config instance"""
    global _config
    if _config is None:
        _config = SecureConfig()
    return _config