File size: 7,326 Bytes
b30f068
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
"""

API Key management for the Agentic AI System.

Loads API keys from configuration file with fallback to defaults.

"""

import json
import os
import secrets
from pathlib import Path
from typing import Dict, Any, List
from src.utils.logging_config import logger


class APIKeyManager:
    """Manages API keys and authentication configuration."""
    
    def __init__(self, config_file: str = "config/api_keys.json"):
        """Initialize API key manager."""
        self.config_file = Path(config_file)
        self.api_keys = {}
        self.settings = {}
        self.load_config()
    
    def load_config(self) -> None:
        """Load API keys from configuration file."""
        try:
            if self.config_file.exists():
                with open(self.config_file, 'r') as f:
                    config = json.load(f)
                
                self.api_keys = config.get('api_keys', {})
                self.settings = config.get('settings', {})
                
                # Filter enabled keys only
                self.api_keys = {
                    key: value for key, value in self.api_keys.items()
                    if value.get('enabled', True)
                }
                
                logger.info(f"Loaded {len(self.api_keys)} API keys from {self.config_file}")
            else:
                logger.warning(f"API keys file not found: {self.config_file}")
                self._load_default_keys()
                
        except Exception as e:
            logger.error(f"Failed to load API keys: {e}")
            self._load_default_keys()
    
    def _load_default_keys(self) -> None:
        """No config file found: generate ONE random read/write key and persist it.



        Security: never ship static, published tokens. A key can also be supplied

        out-of-band via the ADMIN_API_KEY env var (useful for ephemeral/HF-style

        deploys where writing a config file is undesirable). Admin permission is

        never granted by default — add an admin key explicitly with manage_api_keys.

        """
        env_key = os.getenv("ADMIN_API_KEY")
        generated = env_key or secrets.token_urlsafe(32)
        self.api_keys = {
            generated: {
                "name": "Auto-generated key",
                "permissions": ["read", "write"],
                "description": "Generated at first start (no config/api_keys.json found)",
                "rate_limit": 100,
                "enabled": True,
            }
        }
        self.settings = {
            "default_rate_limit": 100,
            "token_expiry_days": 30,
            "require_https": True,
        }
        if env_key:
            logger.info("Loaded API key from ADMIN_API_KEY env var")
        else:
            # Print ONCE so the operator can retrieve it; then persist so it is stable.
            logger.warning(
                "No config/api_keys.json found. Generated a new API key "
                "(shown once): %s", generated
            )
            try:
                self.save_config()
            except Exception as e:  # non-fatal: still usable in-memory this run
                logger.error("Could not persist generated API key: %s", e)
    
    def get_api_keys(self) -> Dict[str, Any]:
        """Get all enabled API keys."""
        return self.api_keys.copy()
    
    def validate_key(self, key: str) -> bool:
        """Check if API key is valid."""
        return key in self.api_keys and self.api_keys[key].get('enabled', True)
    
    def get_key_info(self, key: str) -> Dict[str, Any]:
        """Get information about a specific API key."""
        return self.api_keys.get(key, {})
    
    def has_permission(self, key: str, permission: str) -> bool:
        """Check if API key has specific permission."""
        if key not in self.api_keys:
            return False
        
        permissions = self.api_keys[key].get('permissions', [])
        return permission in permissions
    
    def get_rate_limit(self, key: str) -> int:
        """Get rate limit for API key."""
        if key not in self.api_keys:
            return self.settings.get('default_rate_limit', 100)
        
        return self.api_keys[key].get('rate_limit', self.settings.get('default_rate_limit', 100))
    
    def add_key(self, key: str, config: Dict[str, Any]) -> bool:
        """Add new API key."""
        try:
            self.api_keys[key] = config
            self.save_config()
            logger.info(f"Added new API key: {config.get('name', 'Unknown')}")
            return True
        except Exception as e:
            logger.error(f"Failed to add API key: {e}")
            return False
    
    def remove_key(self, key: str) -> bool:
        """Remove API key."""
        try:
            if key in self.api_keys:
                removed = self.api_keys.pop(key)
                self.save_config()
                logger.info(f"Removed API key: {removed.get('name', 'Unknown')}")
                return True
            return False
        except Exception as e:
            logger.error(f"Failed to remove API key: {e}")
            return False
    
    def disable_key(self, key: str) -> bool:
        """Disable API key without removing it."""
        try:
            if key in self.api_keys:
                self.api_keys[key]['enabled'] = False
                self.save_config()
                logger.info(f"Disabled API key: {key}")
                return True
            return False
        except Exception as e:
            logger.error(f"Failed to disable API key: {e}")
            return False
    
    def save_config(self) -> bool:
        """Save current configuration to file."""
        try:
            # Ensure config directory exists
            self.config_file.parent.mkdir(parents=True, exist_ok=True)
            
            config = {
                "api_keys": self.api_keys,
                "settings": self.settings
            }
            
            with open(self.config_file, 'w') as f:
                json.dump(config, f, indent=2)
            
            logger.info(f"Saved API keys configuration to {self.config_file}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to save API keys: {e}")
            return False
    
    def list_keys(self) -> List[Dict[str, Any]]:
        """List all API keys with their info (excluding the actual key)."""
        keys_info = []
        for key, info in self.api_keys.items():
            key_info = {
                "key_id": key[:8] + "..." + key[-4:],  # Masked key
                "name": info.get('name', 'Unknown'),
                "permissions": info.get('permissions', []),
                "rate_limit": info.get('rate_limit', 100),
                "enabled": info.get('enabled', True),
                "description": info.get('description', '')
            }
            keys_info.append(key_info)
        
        return keys_info
    
    def reload(self) -> None:
        """Reload configuration from file."""
        self.load_config()


# Global instance
api_key_manager = APIKeyManager()