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

2

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

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

5"""

6

7 from __future__ import annotations

8

9 import os

10 from typing import TYPE_CHECKING

11

12 from session_buddy.settings import get_llm_api_key, get_settings

13

14 if TYPE_CHECKING:

15 from mcp_common.security import APIKeyValidator

16

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

18 try:

19 from mcp_common.security import APIKeyValidator

20

21 SECURITY_AVAILABLE = True

22 except ImportError:

23 SECURITY_AVAILABLE = False

24

25

26 def 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

70 def _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

88 def _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

106 def _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

118 def _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

136 def 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