mcp-proxy-adapter 6.3.27__py3-none-any.whl → 6.3.29__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- mcp_proxy_adapter/api/app.py +118 -37
- mcp_proxy_adapter/commands/command_registry.py +24 -1
- mcp_proxy_adapter/core/config_validator.py +172 -234
- mcp_proxy_adapter/core/proxy_client.py +163 -640
- mcp_proxy_adapter/core/proxy_registration.py +143 -41
- mcp_proxy_adapter/main.py +58 -18
- mcp_proxy_adapter/version.py +1 -1
- {mcp_proxy_adapter-6.3.27.dist-info → mcp_proxy_adapter-6.3.29.dist-info}/METADATA +1 -1
- {mcp_proxy_adapter-6.3.27.dist-info → mcp_proxy_adapter-6.3.29.dist-info}/RECORD +12 -12
- {mcp_proxy_adapter-6.3.27.dist-info → mcp_proxy_adapter-6.3.29.dist-info}/WHEEL +0 -0
- {mcp_proxy_adapter-6.3.27.dist-info → mcp_proxy_adapter-6.3.29.dist-info}/entry_points.txt +0 -0
- {mcp_proxy_adapter-6.3.27.dist-info → mcp_proxy_adapter-6.3.29.dist-info}/top_level.txt +0 -0
@@ -1,662 +1,185 @@
|
|
1
1
|
"""
|
2
|
-
Proxy Client Module
|
3
|
-
|
4
|
-
This module provides a client for registering with MCP proxy servers
|
5
|
-
using mcp_security_framework for secure authentication and connections.
|
6
|
-
|
7
2
|
Author: Vasiliy Zdanovskiy
|
8
3
|
email: vasilyvz@gmail.com
|
9
|
-
"""
|
10
|
-
|
11
|
-
import asyncio
|
12
|
-
import json
|
13
|
-
import time
|
14
|
-
import ssl
|
15
|
-
from typing import Dict, Any, Optional, Tuple, List
|
16
|
-
from urllib.parse import urljoin, urlparse
|
17
|
-
from pathlib import Path
|
18
|
-
|
19
|
-
import aiohttp
|
20
|
-
from aiohttp import ClientTimeout, TCPConnector
|
21
|
-
|
22
|
-
# Import framework components
|
23
|
-
try:
|
24
|
-
from mcp_security_framework.core.client_security import ClientSecurityManager
|
25
|
-
from mcp_security_framework.schemas.config import ClientSecurityConfig
|
26
|
-
from mcp_security_framework.schemas.models import AuthResult, ValidationResult
|
27
|
-
from mcp_security_framework.utils.crypto_utils import (
|
28
|
-
generate_api_key,
|
29
|
-
create_jwt_token,
|
30
|
-
)
|
31
|
-
from mcp_security_framework.utils.cert_utils import validate_certificate_format
|
32
|
-
|
33
|
-
SECURITY_FRAMEWORK_AVAILABLE = True
|
34
|
-
except ImportError:
|
35
|
-
SECURITY_FRAMEWORK_AVAILABLE = False
|
36
|
-
ClientSecurityManager = None
|
37
|
-
ClientSecurityConfig = None
|
38
|
-
AuthResult = None
|
39
|
-
ValidationResult = None
|
40
|
-
|
41
|
-
from mcp_proxy_adapter.core.logging import logger
|
42
|
-
|
43
|
-
|
44
|
-
class ProxyClientError(Exception):
|
45
|
-
"""Exception raised when proxy client operations fail."""
|
46
|
-
|
47
|
-
pass
|
48
|
-
|
49
|
-
|
50
|
-
class ProxyClient:
|
51
|
-
"""
|
52
|
-
Client for registering with MCP proxy servers.
|
53
|
-
|
54
|
-
Provides secure registration, heartbeat, and discovery functionality
|
55
|
-
using mcp_security_framework for authentication and SSL/TLS.
|
56
|
-
"""
|
57
|
-
|
58
|
-
def __init__(self, config: Dict[str, Any]):
|
59
|
-
"""
|
60
|
-
Initialize proxy client.
|
61
|
-
|
62
|
-
Args:
|
63
|
-
config: Client configuration
|
64
|
-
"""
|
65
|
-
self.config = config
|
66
|
-
# Try both registration and proxy_registration for backward compatibility
|
67
|
-
self.registration_config = config.get(
|
68
|
-
"registration", config.get("proxy_registration", {})
|
69
|
-
)
|
70
|
-
|
71
|
-
# Basic settings
|
72
|
-
self.proxy_url = self.registration_config.get(
|
73
|
-
"proxy_url", self.registration_config.get("server_url")
|
74
|
-
)
|
75
|
-
self.server_id = self.registration_config.get(
|
76
|
-
"server_id",
|
77
|
-
self.registration_config.get("proxy_info", {}).get(
|
78
|
-
"name", "mcp_proxy_adapter"
|
79
|
-
),
|
80
|
-
)
|
81
|
-
self.server_name = self.registration_config.get(
|
82
|
-
"server_name",
|
83
|
-
self.registration_config.get("proxy_info", {}).get(
|
84
|
-
"name", "MCP Proxy Adapter"
|
85
|
-
),
|
86
|
-
)
|
87
|
-
self.description = self.registration_config.get(
|
88
|
-
"description",
|
89
|
-
self.registration_config.get("proxy_info", {}).get("description", ""),
|
90
|
-
)
|
91
|
-
self.version = self.registration_config.get(
|
92
|
-
"version",
|
93
|
-
self.registration_config.get("proxy_info", {}).get("version", "1.0.0"),
|
94
|
-
)
|
95
|
-
|
96
|
-
# Authentication settings
|
97
|
-
self.auth_method = self.registration_config.get("auth_method", "none")
|
98
|
-
self.auth_config = self._get_auth_config()
|
99
|
-
|
100
|
-
# Heartbeat settings
|
101
|
-
heartbeat_config = self.registration_config.get("heartbeat", {})
|
102
|
-
self.heartbeat_interval = heartbeat_config.get("interval", 300)
|
103
|
-
self.heartbeat_timeout = heartbeat_config.get("timeout", 30)
|
104
|
-
self.retry_attempts = heartbeat_config.get("retry_attempts", 3)
|
105
|
-
self.retry_delay = heartbeat_config.get("retry_delay", 60)
|
106
|
-
|
107
|
-
# Auto discovery settings
|
108
|
-
discovery_config = self.registration_config.get("auto_discovery", {})
|
109
|
-
self.discovery_enabled = discovery_config.get("enabled", False)
|
110
|
-
self.discovery_urls = discovery_config.get("discovery_urls", [])
|
111
|
-
self.discovery_interval = discovery_config.get("discovery_interval", 3600)
|
112
|
-
|
113
|
-
# Initialize security manager
|
114
|
-
self.security_manager = self._create_security_manager()
|
115
|
-
|
116
|
-
# State
|
117
|
-
self.registered = False
|
118
|
-
self.server_key: Optional[str] = None
|
119
|
-
self.server_url: Optional[str] = None
|
120
|
-
self.heartbeat_task: Optional[asyncio.Task] = None
|
121
|
-
self.discovery_task: Optional[asyncio.Task] = None
|
122
|
-
|
123
|
-
logger.info("Proxy client initialized with security framework integration")
|
124
|
-
|
125
|
-
def _get_auth_config(self) -> Dict[str, Any]:
|
126
|
-
"""Get authentication configuration based on auth method."""
|
127
|
-
if self.auth_method == "certificate":
|
128
|
-
return self.registration_config.get("certificate", {})
|
129
|
-
elif self.auth_method == "token":
|
130
|
-
return self.registration_config.get("token", {})
|
131
|
-
elif self.auth_method == "api_key":
|
132
|
-
return self.registration_config.get("api_key", {})
|
133
|
-
else:
|
134
|
-
return {}
|
135
|
-
|
136
|
-
def _create_security_manager(self) -> Optional[ClientSecurityManager]:
|
137
|
-
"""Create client security manager."""
|
138
|
-
if not SECURITY_FRAMEWORK_AVAILABLE:
|
139
|
-
logger.warning("mcp_security_framework not available, using basic client")
|
140
|
-
return None
|
141
|
-
|
142
|
-
try:
|
143
|
-
# Create client security configuration
|
144
|
-
client_security_config = self.registration_config.get("client_security", {})
|
145
|
-
|
146
|
-
if not client_security_config.get("enabled", False):
|
147
|
-
logger.info("Client security disabled in configuration")
|
148
|
-
return None
|
149
4
|
|
150
|
-
|
151
|
-
security_config = {
|
152
|
-
"security": {
|
153
|
-
"ssl": {
|
154
|
-
"enabled": client_security_config.get("ssl_enabled", False),
|
155
|
-
"client_cert_file": client_security_config.get(
|
156
|
-
"certificate_auth", {}
|
157
|
-
).get("cert_file"),
|
158
|
-
"client_key_file": client_security_config.get(
|
159
|
-
"certificate_auth", {}
|
160
|
-
).get("key_file"),
|
161
|
-
"ca_cert_file": client_security_config.get(
|
162
|
-
"certificate_auth", {}
|
163
|
-
).get("ca_cert_file"),
|
164
|
-
"verify_mode": "CERT_REQUIRED",
|
165
|
-
"min_tls_version": "TLSv1.2",
|
166
|
-
"check_hostname": True,
|
167
|
-
"check_expiry": True,
|
168
|
-
},
|
169
|
-
"auth": {
|
170
|
-
"enabled": True,
|
171
|
-
"methods": client_security_config.get(
|
172
|
-
"auth_methods", ["api_key"]
|
173
|
-
),
|
174
|
-
"api_keys": {
|
175
|
-
client_security_config.get("api_key_auth", {}).get(
|
176
|
-
"key", "default"
|
177
|
-
): {
|
178
|
-
"roles": ["proxy_client"],
|
179
|
-
"permissions": ["register", "heartbeat", "discover"],
|
180
|
-
}
|
181
|
-
},
|
182
|
-
},
|
183
|
-
}
|
184
|
-
}
|
5
|
+
Core mTLS Proxy Client.
|
185
6
|
|
186
|
-
|
7
|
+
Provides an asynchronous client for communicating with a proxy-like server over
|
8
|
+
mutual TLS. Designed to be used by services built on this framework to:
|
9
|
+
- perform health/heartbeat checks
|
10
|
+
- register themselves with the proxy
|
187
11
|
|
188
|
-
|
189
|
-
|
190
|
-
|
191
|
-
|
192
|
-
def set_server_url(self, server_url: str) -> None:
|
193
|
-
"""
|
194
|
-
Set the server URL for registration.
|
195
|
-
|
196
|
-
Args:
|
197
|
-
server_url: The URL where this server is accessible.
|
198
|
-
"""
|
199
|
-
self.server_url = server_url
|
200
|
-
logger.info(f"Proxy client server URL set to: {server_url}")
|
201
|
-
|
202
|
-
def _get_auth_headers(self) -> Dict[str, str]:
|
203
|
-
"""
|
204
|
-
Get authentication headers for requests.
|
205
|
-
|
206
|
-
Returns:
|
207
|
-
Dictionary of authentication headers
|
208
|
-
"""
|
209
|
-
headers = {"Content-Type": "application/json"}
|
210
|
-
|
211
|
-
if not self.security_manager:
|
212
|
-
return headers
|
213
|
-
|
214
|
-
try:
|
215
|
-
if self.auth_method == "certificate":
|
216
|
-
return self.security_manager.get_client_auth_headers("certificate")
|
217
|
-
elif self.auth_method == "token":
|
218
|
-
token = self.auth_config.get("token")
|
219
|
-
return self.security_manager.get_client_auth_headers("jwt", token=token)
|
220
|
-
elif self.auth_method == "api_key":
|
221
|
-
api_key = self.auth_config.get("key")
|
222
|
-
return self.security_manager.get_client_auth_headers(
|
223
|
-
"api_key", api_key=api_key
|
224
|
-
)
|
225
|
-
else:
|
226
|
-
return headers
|
227
|
-
except Exception as e:
|
228
|
-
logger.error(f"Failed to get auth headers: {e}")
|
229
|
-
return headers
|
230
|
-
|
231
|
-
def _create_ssl_context(self) -> Optional[ssl.SSLContext]:
|
232
|
-
"""
|
233
|
-
Create SSL context for secure connections.
|
12
|
+
This client intentionally avoids framework-specific configuration objects and
|
13
|
+
accepts explicit parameters for clarity and portability.
|
14
|
+
"""
|
15
|
+
from __future__ import annotations
|
234
16
|
|
235
|
-
|
236
|
-
|
237
|
-
|
238
|
-
|
239
|
-
|
17
|
+
import json
|
18
|
+
import ssl
|
19
|
+
from dataclasses import dataclass
|
20
|
+
from typing import Any, Dict, Optional, Tuple
|
21
|
+
from urllib.parse import urljoin
|
240
22
|
|
241
|
-
|
242
|
-
return self.security_manager.create_client_ssl_context()
|
243
|
-
except Exception as e:
|
244
|
-
logger.error(f"Failed to create SSL context: {e}")
|
245
|
-
return None
|
23
|
+
import aiohttp # type: ignore[import]
|
246
24
|
|
247
|
-
async def register(self) -> bool:
|
248
|
-
"""
|
249
|
-
Register with the proxy server.
|
250
25
|
|
251
|
-
|
252
|
-
|
253
|
-
|
254
|
-
if not self.proxy_url:
|
255
|
-
logger.error("Proxy URL not configured")
|
256
|
-
return False
|
26
|
+
@dataclass(frozen=True)
|
27
|
+
class RegistrationRequest:
|
28
|
+
"""Data payload for registration calls to the proxy server."""
|
257
29
|
|
258
|
-
|
259
|
-
|
260
|
-
|
30
|
+
server_id: str
|
31
|
+
server_name: str
|
32
|
+
description: Optional[str] = None
|
33
|
+
extra: Optional[Dict[str, Any]] = None
|
261
34
|
|
262
|
-
|
263
|
-
|
264
|
-
registration_data = {
|
35
|
+
def to_dict(self) -> Dict[str, Any]:
|
36
|
+
payload: Dict[str, Any] = {
|
265
37
|
"server_id": self.server_id,
|
266
|
-
"server_url": self.server_url,
|
267
38
|
"server_name": self.server_name,
|
268
|
-
"description": self.description,
|
269
|
-
"version": self.version,
|
270
|
-
"capabilities": proxy_info.get("capabilities", ["jsonrpc", "rest"]),
|
271
|
-
"endpoints": proxy_info.get(
|
272
|
-
"endpoints",
|
273
|
-
{"jsonrpc": "/api/jsonrpc", "rest": "/cmd", "health": "/health"},
|
274
|
-
),
|
275
|
-
"auth_method": self.auth_method,
|
276
|
-
"security_enabled": self.security_manager is not None,
|
277
39
|
}
|
40
|
+
if self.description is not None:
|
41
|
+
payload["description"] = self.description
|
42
|
+
if self.extra:
|
43
|
+
payload.update(self.extra)
|
44
|
+
return payload
|
278
45
|
|
279
|
-
logger.info(f"Attempting to register with proxy at {self.proxy_url}")
|
280
|
-
logger.debug(f"Registration data: {registration_data}")
|
281
|
-
|
282
|
-
for attempt in range(self.retry_attempts):
|
283
|
-
try:
|
284
|
-
success, result = await self._make_request(
|
285
|
-
"/register", registration_data
|
286
|
-
)
|
287
|
-
|
288
|
-
if success:
|
289
|
-
self.registered = True
|
290
|
-
self.server_key = result.get("server_key")
|
291
|
-
logger.info(
|
292
|
-
f"✅ Successfully registered with proxy. Server key: {self.server_key}"
|
293
|
-
)
|
294
|
-
|
295
|
-
# Start heartbeat and discovery
|
296
|
-
await self._start_background_tasks()
|
297
|
-
|
298
|
-
return True
|
299
|
-
else:
|
300
|
-
error_msg = result.get("error", {}).get("message", "Unknown error")
|
301
|
-
logger.warning(
|
302
|
-
f"❌ Registration attempt {attempt + 1} failed: {error_msg}"
|
303
|
-
)
|
304
|
-
|
305
|
-
if attempt < self.retry_attempts - 1:
|
306
|
-
logger.info(f"Retrying in {self.retry_delay} seconds...")
|
307
|
-
await asyncio.sleep(self.retry_delay)
|
308
46
|
|
309
|
-
|
310
|
-
|
311
|
-
|
312
|
-
|
313
|
-
|
314
|
-
|
315
|
-
|
316
|
-
|
317
|
-
|
318
|
-
|
319
|
-
|
320
|
-
|
321
|
-
|
322
|
-
|
323
|
-
async def unregister(self) -> bool:
|
324
|
-
"""
|
325
|
-
Unregister from the proxy server.
|
326
|
-
|
327
|
-
Returns:
|
328
|
-
True if unregistration was successful, False otherwise.
|
329
|
-
"""
|
330
|
-
if not self.registered or not self.server_key:
|
331
|
-
logger.info("Not registered with proxy, skipping unregistration")
|
332
|
-
return True
|
333
|
-
|
334
|
-
# Stop background tasks
|
335
|
-
await self._stop_background_tasks()
|
336
|
-
|
337
|
-
# Extract copy_number from server_key
|
338
|
-
try:
|
339
|
-
copy_number = int(self.server_key.split("_")[-1])
|
340
|
-
except (ValueError, IndexError):
|
341
|
-
copy_number = 1
|
342
|
-
|
343
|
-
unregistration_data = {"server_id": self.server_id, "copy_number": copy_number}
|
344
|
-
|
345
|
-
logger.info(f"Attempting to unregister from proxy at {self.proxy_url}")
|
346
|
-
|
347
|
-
try:
|
348
|
-
success, result = await self._make_request(
|
349
|
-
"/unregister", unregistration_data
|
47
|
+
class ProxyClient:
|
48
|
+
"""Asynchronous mTLS HTTP client for communicating with a proxy server.
|
49
|
+
|
50
|
+
Usage:
|
51
|
+
async with ProxyClient(
|
52
|
+
base_url="https://127.0.0.1:3004",
|
53
|
+
ca_cert_path="/path/to/ca.crt",
|
54
|
+
client_cert_path="/path/to/client.crt",
|
55
|
+
client_key_path="/path/to/client.key",
|
56
|
+
) as client:
|
57
|
+
status, health = await client.health()
|
58
|
+
status, hb = await client.heartbeat()
|
59
|
+
status, reg = await client.register(
|
60
|
+
RegistrationRequest(...)
|
350
61
|
)
|
62
|
+
"""
|
351
63
|
|
352
|
-
|
353
|
-
|
354
|
-
|
355
|
-
|
356
|
-
|
357
|
-
|
358
|
-
|
359
|
-
|
360
|
-
|
361
|
-
|
362
|
-
|
363
|
-
|
364
|
-
|
365
|
-
|
366
|
-
|
367
|
-
|
368
|
-
|
369
|
-
|
370
|
-
|
371
|
-
|
372
|
-
|
373
|
-
|
374
|
-
|
375
|
-
|
376
|
-
|
377
|
-
|
378
|
-
|
379
|
-
|
380
|
-
|
381
|
-
|
382
|
-
|
383
|
-
|
384
|
-
|
385
|
-
|
386
|
-
|
387
|
-
|
388
|
-
|
389
|
-
|
390
|
-
|
391
|
-
|
392
|
-
|
393
|
-
|
394
|
-
|
395
|
-
|
396
|
-
|
397
|
-
|
398
|
-
|
399
|
-
|
400
|
-
|
401
|
-
|
402
|
-
|
403
|
-
|
404
|
-
|
405
|
-
|
406
|
-
|
407
|
-
|
408
|
-
|
409
|
-
|
410
|
-
|
411
|
-
|
412
|
-
|
413
|
-
|
414
|
-
|
415
|
-
|
416
|
-
for discovery_url in self.discovery_urls:
|
64
|
+
def __init__(
|
65
|
+
self,
|
66
|
+
base_url: str,
|
67
|
+
*,
|
68
|
+
ca_cert_path: str,
|
69
|
+
client_cert_path: str,
|
70
|
+
client_key_path: str,
|
71
|
+
request_timeout_s: float = 5.0,
|
72
|
+
min_tls_version: ssl.TLSVersion = ssl.TLSVersion.TLSv1_2,
|
73
|
+
verify_mode: ssl.VerifyMode = ssl.CERT_REQUIRED,
|
74
|
+
) -> None:
|
75
|
+
if not base_url.startswith("http"):
|
76
|
+
raise ValueError("base_url must start with http/https")
|
77
|
+
self._base_url: str = base_url.rstrip("/")
|
78
|
+
self._ca_cert_path: str = ca_cert_path
|
79
|
+
self._client_cert_path: str = client_cert_path
|
80
|
+
self._client_key_path: str = client_key_path
|
81
|
+
self._request_timeout_s: float = request_timeout_s
|
82
|
+
self._min_tls_version: ssl.TLSVersion = min_tls_version
|
83
|
+
self._verify_mode: ssl.VerifyMode = verify_mode
|
84
|
+
|
85
|
+
self._ssl_context: Optional[ssl.SSLContext] = None
|
86
|
+
self._session: Optional[aiohttp.ClientSession] = None
|
87
|
+
|
88
|
+
async def __aenter__(self) -> "ProxyClient":
|
89
|
+
await self._ensure_session()
|
90
|
+
return self
|
91
|
+
|
92
|
+
async def __aexit__(self, exc_type, exc, tb) -> None:
|
93
|
+
await self.close()
|
94
|
+
|
95
|
+
async def _ensure_session(self) -> None:
|
96
|
+
if self._session is not None:
|
97
|
+
return
|
98
|
+
ssl_context = self._build_ssl_context()
|
99
|
+
timeout = aiohttp.ClientTimeout(total=self._request_timeout_s)
|
100
|
+
connector = aiohttp.TCPConnector(ssl=ssl_context)
|
101
|
+
self._session = aiohttp.ClientSession(
|
102
|
+
timeout=timeout,
|
103
|
+
connector=connector,
|
104
|
+
)
|
105
|
+
self._ssl_context = ssl_context
|
106
|
+
|
107
|
+
async def close(self) -> None:
|
108
|
+
if self._session is not None:
|
109
|
+
await self._session.close()
|
110
|
+
self._session = None
|
111
|
+
self._ssl_context = None
|
112
|
+
|
113
|
+
def _build_ssl_context(self) -> ssl.SSLContext:
|
114
|
+
ctx = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
|
115
|
+
ctx.minimum_version = self._min_tls_version
|
116
|
+
ctx.verify_mode = self._verify_mode
|
117
|
+
ctx.load_verify_locations(self._ca_cert_path)
|
118
|
+
ctx.load_cert_chain(self._client_cert_path, self._client_key_path)
|
119
|
+
return ctx
|
120
|
+
|
121
|
+
async def _get_json(self, path: str) -> Tuple[int, Dict[str, Any]]:
|
122
|
+
await self._ensure_session()
|
123
|
+
assert self._session is not None
|
124
|
+
url = urljoin(self._base_url + "/", path.lstrip("/"))
|
125
|
+
async with self._session.get(url) as resp:
|
126
|
+
status = resp.status
|
127
|
+
body_text = await resp.text()
|
417
128
|
try:
|
418
|
-
|
419
|
-
|
129
|
+
data: Dict[str, Any] = (
|
130
|
+
json.loads(body_text) if body_text else {}
|
420
131
|
)
|
421
|
-
|
422
|
-
|
423
|
-
|
424
|
-
|
425
|
-
|
426
|
-
|
427
|
-
|
428
|
-
|
429
|
-
|
430
|
-
|
431
|
-
|
432
|
-
|
433
|
-
|
434
|
-
|
435
|
-
|
436
|
-
|
437
|
-
|
438
|
-
|
439
|
-
|
440
|
-
|
441
|
-
|
442
|
-
Args:
|
443
|
-
endpoint: API endpoint
|
444
|
-
data: Request data
|
445
|
-
base_url: Base URL (optional, uses self.proxy_url if not provided)
|
446
|
-
|
447
|
-
Returns:
|
448
|
-
Tuple of (success, result)
|
449
|
-
"""
|
450
|
-
url = urljoin(base_url or self.proxy_url, endpoint)
|
451
|
-
|
452
|
-
# Get authentication headers
|
453
|
-
headers = self._get_auth_headers()
|
454
|
-
|
455
|
-
# Create SSL context if needed
|
456
|
-
ssl_context = self._create_ssl_context()
|
457
|
-
|
458
|
-
# Create connector with SSL context
|
459
|
-
connector = None
|
460
|
-
if ssl_context:
|
461
|
-
connector = TCPConnector(ssl=ssl_context)
|
462
|
-
|
463
|
-
try:
|
464
|
-
timeout = ClientTimeout(total=self.heartbeat_timeout)
|
465
|
-
|
466
|
-
async with aiohttp.ClientSession(connector=connector) as session:
|
467
|
-
async with session.post(
|
468
|
-
url, json=data, headers=headers, timeout=timeout
|
469
|
-
) as response:
|
470
|
-
result = await response.json()
|
471
|
-
|
472
|
-
# Validate response if security manager available
|
473
|
-
if self.security_manager:
|
474
|
-
self.security_manager.validate_server_response(
|
475
|
-
dict(response.headers)
|
476
|
-
)
|
477
|
-
|
478
|
-
return response.status == 200, result
|
479
|
-
except Exception as e:
|
480
|
-
logger.error(f"Request failed: {e}")
|
481
|
-
return False, {"error": {"message": str(e)}}
|
482
|
-
finally:
|
483
|
-
if connector:
|
484
|
-
await connector.close()
|
485
|
-
|
486
|
-
async def _start_background_tasks(self) -> None:
|
487
|
-
"""Start heartbeat and discovery background tasks."""
|
488
|
-
# Start heartbeat
|
489
|
-
if self.registration_config.get("heartbeat", {}).get("enabled", True):
|
490
|
-
self.heartbeat_task = asyncio.create_task(self._heartbeat_loop())
|
491
|
-
logger.info("Heartbeat task started")
|
492
|
-
|
493
|
-
# Start discovery
|
494
|
-
if self.discovery_enabled:
|
495
|
-
self.discovery_task = asyncio.create_task(self._discovery_loop())
|
496
|
-
logger.info("Discovery task started")
|
497
|
-
|
498
|
-
async def _stop_background_tasks(self) -> None:
|
499
|
-
"""Stop background tasks."""
|
500
|
-
# Stop heartbeat
|
501
|
-
if self.heartbeat_task and not self.heartbeat_task.done():
|
502
|
-
self.heartbeat_task.cancel()
|
503
|
-
try:
|
504
|
-
await self.heartbeat_task
|
505
|
-
except asyncio.CancelledError:
|
506
|
-
pass
|
507
|
-
logger.info("Heartbeat task stopped")
|
508
|
-
|
509
|
-
# Stop discovery
|
510
|
-
if self.discovery_task and not self.discovery_task.done():
|
511
|
-
self.discovery_task.cancel()
|
512
|
-
try:
|
513
|
-
await self.discovery_task
|
514
|
-
except asyncio.CancelledError:
|
515
|
-
pass
|
516
|
-
logger.info("Discovery task stopped")
|
517
|
-
|
518
|
-
async def _heartbeat_loop(self) -> None:
|
519
|
-
"""Heartbeat loop to keep registration alive."""
|
520
|
-
while self.registered:
|
521
|
-
try:
|
522
|
-
await asyncio.sleep(self.heartbeat_interval)
|
523
|
-
|
524
|
-
if not self.registered:
|
525
|
-
break
|
526
|
-
|
527
|
-
# Send heartbeat
|
528
|
-
success = await self.send_heartbeat()
|
529
|
-
if not success:
|
530
|
-
logger.warning("Heartbeat failed, attempting to re-register")
|
531
|
-
await self.register()
|
532
|
-
|
533
|
-
except asyncio.CancelledError:
|
534
|
-
break
|
535
|
-
except Exception as e:
|
536
|
-
logger.error(f"Heartbeat error: {e}")
|
537
|
-
|
538
|
-
async def _discovery_loop(self) -> None:
|
539
|
-
"""Discovery loop to find new proxy servers."""
|
540
|
-
while self.registered:
|
132
|
+
except json.JSONDecodeError:
|
133
|
+
data = {"raw": body_text}
|
134
|
+
return status, data
|
135
|
+
|
136
|
+
async def _post_json(
|
137
|
+
self,
|
138
|
+
path: str,
|
139
|
+
payload: Dict[str, Any],
|
140
|
+
) -> Tuple[int, Dict[str, Any]]:
|
141
|
+
await self._ensure_session()
|
142
|
+
assert self._session is not None
|
143
|
+
url = urljoin(self._base_url + "/", path.lstrip("/"))
|
144
|
+
headers = {"Content-Type": "application/json"}
|
145
|
+
async with self._session.post(
|
146
|
+
url,
|
147
|
+
headers=headers,
|
148
|
+
json=payload,
|
149
|
+
) as resp:
|
150
|
+
status = resp.status
|
151
|
+
body_text = await resp.text()
|
541
152
|
try:
|
542
|
-
|
543
|
-
|
544
|
-
|
545
|
-
|
546
|
-
|
547
|
-
|
548
|
-
|
549
|
-
|
550
|
-
|
551
|
-
|
552
|
-
|
553
|
-
|
554
|
-
|
555
|
-
|
556
|
-
|
557
|
-
|
558
|
-
|
559
|
-
|
560
|
-
|
561
|
-
|
562
|
-
|
563
|
-
|
564
|
-
|
565
|
-
|
566
|
-
|
567
|
-
|
568
|
-
|
569
|
-
|
570
|
-
|
571
|
-
|
572
|
-
|
573
|
-
|
574
|
-
|
575
|
-
Get current client status.
|
576
|
-
|
577
|
-
Returns:
|
578
|
-
Dictionary with client status information.
|
579
|
-
"""
|
580
|
-
status = {
|
581
|
-
"enabled": self.registration_config.get("enabled", False),
|
582
|
-
"registered": self.registered,
|
583
|
-
"server_key": self.server_key,
|
584
|
-
"server_url": self.server_url,
|
585
|
-
"proxy_url": self.proxy_url,
|
586
|
-
"server_id": self.server_id,
|
587
|
-
"auth_method": self.auth_method,
|
588
|
-
"heartbeat_active": self.heartbeat_task is not None
|
589
|
-
and not self.heartbeat_task.done(),
|
590
|
-
"discovery_active": self.discovery_task is not None
|
591
|
-
and not self.discovery_task.done(),
|
592
|
-
}
|
593
|
-
|
594
|
-
# Add security information
|
595
|
-
if self.security_manager:
|
596
|
-
status["security_enabled"] = True
|
597
|
-
status["ssl_enabled"] = self.security_manager.is_ssl_enabled()
|
598
|
-
status["auth_methods"] = self.security_manager.get_supported_auth_methods()
|
599
|
-
else:
|
600
|
-
status["security_enabled"] = False
|
601
|
-
|
602
|
-
return status
|
603
|
-
|
604
|
-
|
605
|
-
# Global proxy client instance
|
606
|
-
proxy_client: Optional[ProxyClient] = None
|
607
|
-
|
608
|
-
|
609
|
-
def initialize_proxy_client(config: Dict[str, Any]) -> None:
|
610
|
-
"""
|
611
|
-
Initialize global proxy client.
|
612
|
-
|
613
|
-
Args:
|
614
|
-
config: Application configuration
|
615
|
-
"""
|
616
|
-
global proxy_client
|
617
|
-
proxy_client = ProxyClient(config)
|
618
|
-
|
619
|
-
|
620
|
-
async def register_with_proxy(server_url: str) -> bool:
|
621
|
-
"""
|
622
|
-
Register with proxy server.
|
623
|
-
|
624
|
-
Args:
|
625
|
-
server_url: The URL where this server is accessible.
|
626
|
-
|
627
|
-
Returns:
|
628
|
-
True if registration was successful, False otherwise.
|
629
|
-
"""
|
630
|
-
if not proxy_client:
|
631
|
-
logger.error("Proxy client not initialized")
|
632
|
-
return False
|
633
|
-
|
634
|
-
proxy_client.set_server_url(server_url)
|
635
|
-
return await proxy_client.register()
|
636
|
-
|
637
|
-
|
638
|
-
async def unregister_from_proxy() -> bool:
|
639
|
-
"""
|
640
|
-
Unregister from proxy server.
|
641
|
-
|
642
|
-
Returns:
|
643
|
-
True if unregistration was successful, False otherwise.
|
644
|
-
"""
|
645
|
-
if not proxy_client:
|
646
|
-
logger.error("Proxy client not initialized")
|
647
|
-
return False
|
648
|
-
|
649
|
-
return await proxy_client.unregister()
|
650
|
-
|
651
|
-
|
652
|
-
def get_proxy_client_status() -> Dict[str, Any]:
|
653
|
-
"""
|
654
|
-
Get proxy client status.
|
655
|
-
|
656
|
-
Returns:
|
657
|
-
Dictionary with client status information.
|
658
|
-
"""
|
659
|
-
if not proxy_client:
|
660
|
-
return {"error": "Proxy client not initialized"}
|
661
|
-
|
662
|
-
return proxy_client.get_status()
|
153
|
+
data: Dict[str, Any] = (
|
154
|
+
json.loads(body_text) if body_text else {}
|
155
|
+
)
|
156
|
+
except json.JSONDecodeError:
|
157
|
+
data = {"raw": body_text}
|
158
|
+
return status, data
|
159
|
+
|
160
|
+
async def health(
|
161
|
+
self,
|
162
|
+
path: str = "/health",
|
163
|
+
) -> Tuple[int, Dict[str, Any]]:
|
164
|
+
"""Perform a health check against the proxy."""
|
165
|
+
return await self._get_json(path)
|
166
|
+
|
167
|
+
async def heartbeat(
|
168
|
+
self,
|
169
|
+
path: str = "/heartbeat",
|
170
|
+
) -> Tuple[int, Dict[str, Any]]:
|
171
|
+
"""Perform a heartbeat (liveness) check against the proxy."""
|
172
|
+
return await self._get_json(path)
|
173
|
+
|
174
|
+
async def register(
|
175
|
+
self,
|
176
|
+
request: RegistrationRequest,
|
177
|
+
*,
|
178
|
+
path: str = "/register",
|
179
|
+
) -> Tuple[int, Dict[str, Any]]:
|
180
|
+
"""Register the current service on the proxy server."""
|
181
|
+
payload = request.to_dict()
|
182
|
+
return await self._post_json(
|
183
|
+
path,
|
184
|
+
payload,
|
185
|
+
)
|