Coverage for session_buddy / llm / security.py: 9.38%

1"""LLM provider API key security and validation utilities. 

2 

3This module provides security utilities for validating and masking LLM provider 

4API keys during server startup (Phase 3 Security Hardening). 

5""" 

6 

7from __future__ import annotations 

8 

9import os 

10from typing import TYPE_CHECKING 

11 

12from session_buddy.settings import get_llm_api_key, get_settings 

13 

14if TYPE_CHECKING: 

15 from mcp_common.security import APIKeyValidator 

16 

17# Import mcp-common security utilities for API key validation (Phase 3 Security Hardening) 

18try: 

19 from mcp_common.security import APIKeyValidator 

20 

21 SECURITY_AVAILABLE = True 

22except ImportError: 

23 SECURITY_AVAILABLE = False 

24 

25 

26def get_masked_api_key(provider: str = "openai") -> str: 

27 """Get masked API key for safe logging. 

28 

29 Args: 

30 provider: Provider name ('openai', 'gemini', 'ollama') 

31 

32 Returns: 

33 Masked API key string (e.g., "sk-...abc1") for safe display in logs 

34 

35 """ 

36 settings = get_settings() 

37 key_field_map = { 

38 "openai": "openai_api_key", 

39 "anthropic": "anthropic_api_key", 

40 "gemini": "gemini_api_key", 

41 } 

42 key_field = key_field_map.get(provider) 

43 if key_field: 

44 configured = getattr(settings, key_field, None) 

45 if isinstance(configured, str) and configured.strip(): 

46 return settings.get_masked_key(key_name=key_field, visible_chars=4) 

47 

48 api_key = None 

49 

50 if provider == "openai": 

51 api_key = os.getenv("OPENAI_API_KEY") 

52 elif provider == "gemini": 

53 api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY") 

54 elif provider == "ollama": 

55 # Ollama is local, no API key needed 

56 return "N/A (local service)" 

57 

58 if not api_key: 

59 return "***" 

60 

61 if SECURITY_AVAILABLE: 

62 return APIKeyValidator.mask_key(api_key, visible_chars=4) 

63 

64 # Fallback masking without security module 

65 if len(api_key) <= 4: 

66 return "***" 

67 return f"...{api_key[-4:]}" 

68 

69 

70def _get_provider_api_key_and_env(provider: str) -> tuple[str | None, str | None]: 

71 """Get API key and environment variable name for provider.""" 

72 configured_key = get_llm_api_key(provider) 

73 if configured_key: 

74 return configured_key, f"settings.{provider}_api_key" 

75 if provider == "openai": 

76 return os.getenv("OPENAI_API_KEY"), "OPENAI_API_KEY" 

77 if provider == "anthropic": 

78 return os.getenv("ANTHROPIC_API_KEY"), "ANTHROPIC_API_KEY" 

79 if provider == "gemini": 

80 api_key = os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY") 

81 env_var_name = ( 

82 "GEMINI_API_KEY" if os.getenv("GEMINI_API_KEY") else "GOOGLE_API_KEY" 

83 ) 

84 return api_key, env_var_name 

85 return None, None 

86 

87 

88def _validate_provider_with_security(provider: str, api_key: str) -> tuple[bool, str]: 

89 """Validate provider API key using mcp-common security module. 

90 

91 Returns: 

92 Tuple of (success, status_message) 

93 

94 """ 

95 import sys 

96 

97 validator = APIKeyValidator(provider=provider) 

98 try: 

99 validator.validate(api_key, raise_on_invalid=True) 

100 get_masked_api_key(provider) 

101 return True, "valid" 

102 except ValueError: 

103 sys.exit(1) 

104 

105 

106def _validate_provider_basic(provider: str, api_key: str) -> str: 

107 """Basic API key validation without security module. 

108 

109 Returns: 

110 Status message 

111 

112 """ 

113 if len(api_key) < 16: 

114 pass 

115 return "basic_check" 

116 

117 

118def _get_configured_providers() -> list[str]: 

119 """Get list of configured LLM providers.""" 

120 providers: set[str] = set() 

121 if get_llm_api_key("openai"): 

122 providers.add("openai") 

123 if get_llm_api_key("gemini"): 

124 providers.add("gemini") 

125 if get_llm_api_key("anthropic"): 

126 providers.add("anthropic") 

127 if os.getenv("OPENAI_API_KEY"): 

128 providers.add("openai") 

129 if os.getenv("ANTHROPIC_API_KEY"): 

130 providers.add("anthropic") 

131 if os.getenv("GEMINI_API_KEY") or os.getenv("GOOGLE_API_KEY"): 

132 providers.add("gemini") 

133 return sorted(providers) 

134 

135 

136def validate_llm_api_keys_at_startup() -> dict[str, str]: 

137 """Validate LLM provider API keys at server startup (Phase 3 Security Hardening). 

138 

139 Validates API keys for all configured LLM providers (OpenAI, Gemini). 

140 Ollama is skipped as it's a local service without API key requirements. 

141 

142 Returns: 

143 Dictionary mapping provider names to validation status messages 

144 

145 Raises: 

146 SystemExit: If required API keys are invalid or missing 

147 

148 """ 

149 import sys 

150 

151 validated_providers: dict[str, str] = {} 

152 providers_configured = _get_configured_providers() 

153 

154 # If no providers configured, warn but allow startup (Ollama might be used) 

155 if not providers_configured: 

156 return validated_providers 

157 

158 # Validate each configured provider 

159 for provider in providers_configured: 

160 api_key, _env_var_name = _get_provider_api_key_and_env(provider) 

161 

162 if not api_key or not api_key.strip(): 

163 sys.exit(1) 

164 

165 if SECURITY_AVAILABLE: 

166 _, status = _validate_provider_with_security(provider, api_key) 

167 validated_providers[provider] = status 

168 else: 

169 status = _validate_provider_basic(provider, api_key) 

170 validated_providers[provider] = status 

171 

172 return validated_providers