Coverage for me2ai_mcp\auth.py: 0%

93 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-04-13 11:30 +0200

1""" 

2Authentication utilities for ME2AI MCP servers. 

3 

4This module provides authentication managers and methods for securing 

5MCP server endpoints and API access. 

6""" 

7from typing import Dict, List, Any, Optional, Union, Callable, Protocol 

8import os 

9import logging 

10import json 

11from abc import ABC, abstractmethod 

12from dataclasses import dataclass 

13from pathlib import Path 

14from dotenv import load_dotenv 

15 

16# Configure logging 

17logger = logging.getLogger("me2ai-mcp-auth") 

18 

19 

20class AuthProvider(ABC): 

21 """Abstract base class for authentication providers.""" 

22 

23 @abstractmethod 

24 def authenticate(self, credentials: Dict[str, Any]) -> bool: 

25 """Authenticate a request using the provided credentials. 

26  

27 Args: 

28 credentials: Authentication credentials 

29  

30 Returns: 

31 Whether authentication was successful 

32 """ 

33 pass 

34 

35 @abstractmethod 

36 def get_auth_headers(self) -> Dict[str, str]: 

37 """Get authentication headers for outgoing requests. 

38  

39 Returns: 

40 Dictionary of authentication headers 

41 """ 

42 pass 

43 

44 

45class APIKeyAuth(AuthProvider): 

46 """API key-based authentication for MCP servers.""" 

47 

48 def __init__( 

49 self, 

50 api_key: Optional[str] = None, 

51 env_var_name: Optional[str] = None, 

52 header_name: str = "X-API-Key" 

53 ) -> None: 

54 """Initialize API key authentication. 

55  

56 Args: 

57 api_key: API key (optional if env_var_name is provided) 

58 env_var_name: Name of environment variable containing API key 

59 header_name: Name of header for API key 

60 """ 

61 self.header_name = header_name 

62 

63 # Load from environment if not provided directly 

64 if api_key is None and env_var_name: 

65 load_dotenv() 

66 api_key = os.getenv(env_var_name) 

67 

68 self.api_key = api_key 

69 

70 if not self.api_key: 

71 logger.warning(f"No API key provided for {self.__class__.__name__}") 

72 

73 def authenticate(self, credentials: Dict[str, Any]) -> bool: 

74 """Authenticate using API key. 

75  

76 Args: 

77 credentials: Dictionary containing the API key 

78  

79 Returns: 

80 Whether authentication was successful 

81 """ 

82 if not self.api_key: 

83 # If no API key is configured, authentication is disabled 

84 return True 

85 

86 # Extract API key from various potential sources 

87 request_api_key = credentials.get(self.header_name, credentials.get("api_key")) 

88 

89 if not request_api_key: 

90 logger.warning("Authentication failed: No API key provided in request") 

91 return False 

92 

93 # Compare API keys 

94 return request_api_key == self.api_key 

95 

96 def get_auth_headers(self) -> Dict[str, str]: 

97 """Get API key authentication headers. 

98  

99 Returns: 

100 Dictionary containing API key header 

101 """ 

102 if not self.api_key: 

103 return {} 

104 

105 return {self.header_name: self.api_key} 

106 

107 

108class TokenAuth(AuthProvider): 

109 """Token-based authentication for MCP servers.""" 

110 

111 def __init__( 

112 self, 

113 token: Optional[str] = None, 

114 env_var_name: Optional[str] = None, 

115 auth_scheme: str = "Bearer" 

116 ) -> None: 

117 """Initialize token authentication. 

118  

119 Args: 

120 token: Authentication token (optional if env_var_name is provided) 

121 env_var_name: Name of environment variable containing token 

122 auth_scheme: Authentication scheme (e.g., "Bearer") 

123 """ 

124 self.auth_scheme = auth_scheme 

125 

126 # Load from environment if not provided directly 

127 if token is None and env_var_name: 

128 load_dotenv() 

129 token = os.getenv(env_var_name) 

130 

131 self.token = token 

132 

133 if not self.token: 

134 logger.warning(f"No token provided for {self.__class__.__name__}") 

135 

136 def authenticate(self, credentials: Dict[str, Any]) -> bool: 

137 """Authenticate using token. 

138  

139 Args: 

140 credentials: Dictionary containing the authorization header 

141  

142 Returns: 

143 Whether authentication was successful 

144 """ 

145 if not self.token: 

146 # If no token is configured, authentication is disabled 

147 return True 

148 

149 # Extract token from authorization header 

150 auth_header = credentials.get("Authorization", "") 

151 

152 if not auth_header.startswith(f"{self.auth_scheme} "): 

153 logger.warning(f"Authentication failed: Invalid Authorization header format (expected {self.auth_scheme})") 

154 return False 

155 

156 # Extract the token part 

157 request_token = auth_header[len(f"{self.auth_scheme} "):] 

158 

159 # Compare tokens 

160 return request_token == self.token 

161 

162 def get_auth_headers(self) -> Dict[str, str]: 

163 """Get token authentication headers. 

164  

165 Returns: 

166 Dictionary containing Authorization header 

167 """ 

168 if not self.token: 

169 return {} 

170 

171 return {"Authorization": f"{self.auth_scheme} {self.token}"} 

172 

173 

174class AuthManager: 

175 """Authentication manager for ME2AI MCP servers.""" 

176 

177 def __init__(self, providers: Optional[List[AuthProvider]] = None) -> None: 

178 """Initialize the authentication manager. 

179  

180 Args: 

181 providers: List of authentication providers 

182 """ 

183 self.providers = providers or [] 

184 self.logger = logging.getLogger("me2ai-mcp-auth-manager") 

185 

186 def add_provider(self, provider: AuthProvider) -> None: 

187 """Add an authentication provider. 

188  

189 Args: 

190 provider: Authentication provider to add 

191 """ 

192 self.providers.append(provider) 

193 

194 def authenticate(self, credentials: Dict[str, Any]) -> bool: 

195 """Authenticate a request using all providers. 

196  

197 Authentication succeeds if ANY provider authenticates successfully. 

198 If no providers are configured, authentication is always successful. 

199  

200 Args: 

201 credentials: Authentication credentials 

202  

203 Returns: 

204 Whether authentication was successful 

205 """ 

206 if not self.providers: 

207 # If no providers are configured, authentication is disabled 

208 return True 

209 

210 # Try each provider 

211 for provider in self.providers: 

212 if provider.authenticate(credentials): 

213 return True 

214 

215 self.logger.warning("Authentication failed: No provider authenticated the request") 

216 return False 

217 

218 def get_auth_headers(self) -> Dict[str, str]: 

219 """Get authentication headers from the first provider. 

220  

221 Returns: 

222 Dictionary of authentication headers 

223 """ 

224 if not self.providers: 

225 return {} 

226 

227 # Use the first provider's headers 

228 return self.providers[0].get_auth_headers() 

229 

230 @classmethod 

231 def from_env(cls, *env_var_names: str) -> "AuthManager": 

232 """Create an authentication manager from environment variables. 

233  

234 This method creates API key authentication providers for each 

235 environment variable name provided. 

236  

237 Args: 

238 env_var_names: Names of environment variables containing API keys 

239  

240 Returns: 

241 Configured authentication manager 

242 """ 

243 load_dotenv() 

244 

245 providers = [] 

246 

247 for env_var_name in env_var_names: 

248 if os.getenv(env_var_name): 

249 providers.append(APIKeyAuth(env_var_name=env_var_name)) 

250 

251 return cls(providers) 

252 

253 @classmethod 

254 def from_github_token(cls) -> "AuthManager": 

255 """Create an authentication manager using a GitHub token. 

256  

257 This method checks the following environment variables in order: 

258 - GITHUB_API_KEY 

259 - GITHUB_TOKEN 

260 - GITHUB_ACCESS_TOKEN 

261  

262 Returns: 

263 Configured authentication manager with GitHub token 

264 """ 

265 load_dotenv() 

266 

267 # Try different potential environment variable names 

268 token = ( 

269 os.getenv("GITHUB_API_KEY") or 

270 os.getenv("GITHUB_TOKEN") or 

271 os.getenv("GITHUB_ACCESS_TOKEN") 

272 ) 

273 

274 if token: 

275 return cls([TokenAuth(token=token, auth_scheme="Bearer")]) 

276 else: 

277 logger.warning("No GitHub token found in environment variables") 

278 return cls()