lange-python 0.3.16__tar.gz → 0.3.19__tar.gz

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.
Files changed (32) hide show
  1. {lange_python-0.3.16 → lange_python-0.3.19}/PKG-INFO +5 -1
  2. {lange_python-0.3.16 → lange_python-0.3.19}/README.md +4 -0
  3. lange_python-0.3.19/lange/tunnel/_client.py +685 -0
  4. {lange_python-0.3.16 → lange_python-0.3.19}/pyproject.toml +1 -1
  5. lange_python-0.3.16/lange/tunnel/_client.py +0 -477
  6. {lange_python-0.3.16 → lange_python-0.3.19}/lange/__init__.py +0 -0
  7. {lange_python-0.3.16 → lange_python-0.3.19}/lange/__main__.py +0 -0
  8. {lange_python-0.3.16 → lange_python-0.3.19}/lange/_util/__init__.py +0 -0
  9. {lange_python-0.3.16 → lange_python-0.3.19}/lange/_util/_base_client.py +0 -0
  10. {lange_python-0.3.16 → lange_python-0.3.19}/lange/_util/_key_handling.py +0 -0
  11. {lange_python-0.3.16 → lange_python-0.3.19}/lange/cli/__init__.py +0 -0
  12. {lange_python-0.3.16 → lange_python-0.3.19}/lange/cli/build/__init__.py +0 -0
  13. {lange_python-0.3.16 → lange_python-0.3.19}/lange/cli/build/_command.py +0 -0
  14. {lange_python-0.3.16 → lange_python-0.3.19}/lange/cli/build/_discovery.py +0 -0
  15. {lange_python-0.3.16 → lange_python-0.3.19}/lange/cli/build/_docker.py +0 -0
  16. {lange_python-0.3.16 → lange_python-0.3.19}/lange/cli/build/_poetry.py +0 -0
  17. {lange_python-0.3.16 → lange_python-0.3.19}/lange/cli/build/_types.py +0 -0
  18. {lange_python-0.3.16 → lange_python-0.3.19}/lange/cli/code/__init__.py +0 -0
  19. {lange_python-0.3.16 → lange_python-0.3.19}/lange/cli/code/_stats.py +0 -0
  20. {lange_python-0.3.16 → lange_python-0.3.19}/lange/cli/code/audit/__init__.py +0 -0
  21. {lange_python-0.3.16 → lange_python-0.3.19}/lange/cli/code/audit/_command.py +0 -0
  22. {lange_python-0.3.16 → lange_python-0.3.19}/lange/cli/code/audit/_discovery.py +0 -0
  23. {lange_python-0.3.16 → lange_python-0.3.19}/lange/cli/code/audit/_runner.py +0 -0
  24. {lange_python-0.3.16 → lange_python-0.3.19}/lange/cli/code/audit/_types.py +0 -0
  25. {lange_python-0.3.16 → lange_python-0.3.19}/lange/cli/distribution/__init__.py +0 -0
  26. {lange_python-0.3.16 → lange_python-0.3.19}/lange/cli/distribution/_command.py +0 -0
  27. {lange_python-0.3.16 → lange_python-0.3.19}/lange/distribution/__init__.py +0 -0
  28. {lange_python-0.3.16 → lange_python-0.3.19}/lange/distribution/_client.py +0 -0
  29. {lange_python-0.3.16 → lange_python-0.3.19}/lange/distribution/_update_macos.py +0 -0
  30. {lange_python-0.3.16 → lange_python-0.3.19}/lange/distribution/_util.py +0 -0
  31. {lange_python-0.3.16 → lange_python-0.3.19}/lange/tunnel/__init__.py +0 -0
  32. {lange_python-0.3.16 → lange_python-0.3.19}/lange/tunnel/_util.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: lange-python
3
- Version: 0.3.16
3
+ Version: 0.3.19
4
4
  Summary: A bundeld set of tools, clients for the lange-suite of tools and more.
5
5
  Author: contact@robertlange.me
6
6
  Requires-Python: >=3.10
@@ -89,3 +89,7 @@ tunnel.stop()
89
89
  Tunnel clients also expose a `status` property with one of:
90
90
  `"unauthenticated"`, `"off"`, `"pending"`, `"connected"`, or `"failed"`.
91
91
 
92
+ The worker connects to `/api/v1/tunnels/worker`, consumes the
93
+ `worker_connected` welcome payload, forwards HTTP bodies as base64-encoded
94
+ bytes, and supports websocket relay sessions through the same control socket.
95
+
@@ -70,3 +70,7 @@ tunnel.stop()
70
70
 
71
71
  Tunnel clients also expose a `status` property with one of:
72
72
  `"unauthenticated"`, `"off"`, `"pending"`, `"connected"`, or `"failed"`.
73
+
74
+ The worker connects to `/api/v1/tunnels/worker`, consumes the
75
+ `worker_connected` welcome payload, forwards HTTP bodies as base64-encoded
76
+ bytes, and supports websocket relay sessions through the same control socket.
@@ -0,0 +1,685 @@
1
+ from __future__ import annotations
2
+
3
+ import asyncio
4
+ import base64
5
+ import json
6
+ import logging
7
+ import ssl
8
+ import threading
9
+ from typing import Any, Optional
10
+ from urllib.parse import urljoin, urlparse, urlunparse
11
+
12
+ import httpx
13
+ import websockets
14
+ from httpx import Timeout
15
+
16
+ from .._util import BaseLangeLabsClient
17
+ from .._util._base_client import _UNSET_API_KEY
18
+ from ._util import _filter_hop_by_hop_headers
19
+
20
+ logger = logging.getLogger("lange.tunnel")
21
+ HEARTBEAT_INTERVAL_SECONDS = 10.0
22
+
23
+
24
+ class Tunnel(BaseLangeLabsClient):
25
+ """Thread-based tunnel worker client."""
26
+
27
+ def __init__(
28
+ self,
29
+ host: str = "wss://tunnel.lange-labs.com",
30
+ api_key: str | None = None,
31
+ target: str = "http://localhost:80",
32
+ verify_ssl: bool = True,
33
+ max_retries: int = 5,
34
+ retry_delay: float = 5.0,
35
+ open_timeout: float = 20.0,
36
+ daemon: bool = True,
37
+ ) -> None:
38
+ """Initialize one tunnel worker client."""
39
+ super().__init__(api_key=api_key, daemon=daemon, host=host)
40
+ self.target = target.rstrip("/")
41
+ self.verify_ssl = verify_ssl
42
+ self.max_retries = max_retries
43
+ self.retry_delay = retry_delay
44
+ self.open_timeout = open_timeout
45
+
46
+ self._max_retry_delay = 60.0
47
+ self._stop_event = threading.Event()
48
+ self._reconnect_event = threading.Event()
49
+ self._connected = False
50
+ self._remote_address: Optional[str] = None
51
+ self._remote_address_roundrobin: Optional[str] = None
52
+ self._worker_index = -1
53
+ self._pool_size = 0
54
+ self._reconnect_count = 0
55
+ self._worker_id: Optional[str] = None
56
+ self._instance_id: Optional[str] = None
57
+ self._tunnel_id: Optional[str] = None
58
+ self._tunnel_name: Optional[str] = None
59
+ self._max_workers: int = -1
60
+ self._can_consume_tunnels = False
61
+ self._lock = threading.Lock()
62
+ self._loop: Optional[asyncio.AbstractEventLoop] = None
63
+ self._active_ws: Any = None
64
+ self._ws_sessions: dict[str, Any] = {}
65
+ self._ws_reader_tasks: dict[str, asyncio.Task[None]] = {}
66
+
67
+ @property
68
+ def connected(self) -> bool:
69
+ """Get current connection state."""
70
+ with self._lock:
71
+ return self._connected
72
+
73
+ @property
74
+ def status(self) -> str:
75
+ """Get the current tunnel lifecycle status."""
76
+ with self._lock:
77
+ return self._status
78
+
79
+ @property
80
+ def remote_address(self) -> Optional[str]:
81
+ """Get worker-specific public address when provided by the server."""
82
+ with self._lock:
83
+ return self._remote_address
84
+
85
+ @property
86
+ def remote_address_roundrobin(self) -> Optional[str]:
87
+ """Get round-robin public address when provided by the server."""
88
+ with self._lock:
89
+ return self._remote_address_roundrobin
90
+
91
+ @property
92
+ def worker_index(self) -> int:
93
+ """Get current worker index from the latest welcome payload."""
94
+ with self._lock:
95
+ return self._worker_index
96
+
97
+ @property
98
+ def pool_size(self) -> int:
99
+ """Get current worker pool size from the latest welcome payload."""
100
+ with self._lock:
101
+ return self._pool_size
102
+
103
+ @property
104
+ def reconnect_count(self) -> int:
105
+ """Get number of reconnect attempts after the last successful connection."""
106
+ with self._lock:
107
+ return self._reconnect_count
108
+
109
+ @property
110
+ def worker_id(self) -> Optional[str]:
111
+ """Get the worker identifier assigned by the server."""
112
+ with self._lock:
113
+ return self._worker_id
114
+
115
+ @property
116
+ def instance_id(self) -> Optional[str]:
117
+ """Get the API instance identifier from the welcome payload."""
118
+ with self._lock:
119
+ return self._instance_id
120
+
121
+ @property
122
+ def tunnel_name(self) -> Optional[str]:
123
+ """Get the tunnel name from the welcome payload."""
124
+ with self._lock:
125
+ return self._tunnel_name
126
+
127
+ @property
128
+ def max_workers(self) -> int:
129
+ """Get the configured tunnel worker limit from the welcome payload."""
130
+ with self._lock:
131
+ return self._max_workers
132
+
133
+ @property
134
+ def can_consume_tunnels(self) -> bool:
135
+ """Get the worker tunnel-consumption capability from the welcome payload."""
136
+ with self._lock:
137
+ return self._can_consume_tunnels
138
+
139
+ def run(self) -> None:
140
+ """Start the worker loop inside the thread."""
141
+ asyncio.run(self._run_async())
142
+
143
+ def stop(self) -> None:
144
+ """Request a graceful shutdown."""
145
+ self._stop_event.set()
146
+
147
+ def reconnect(self) -> None:
148
+ """Force a full reconnect cycle."""
149
+ self._request_reconnect()
150
+
151
+ def reload(self, api_key: str | None | object = _UNSET_API_KEY) -> None:
152
+ """Reload tunnel authentication and restart the connection flow when running."""
153
+ super().reload(api_key=api_key)
154
+ self._set_connected(False)
155
+
156
+ if self.is_alive() or self._loop is not None:
157
+ self._request_reconnect()
158
+ return
159
+
160
+ self._set_status(self._status_for_auth())
161
+
162
+ async def _run_async(self) -> None:
163
+ """Run the worker connection loop with reconnect backoff."""
164
+ tunnel_url = self._build_tunnel_url()
165
+ ssl_context = self._build_ssl_context(tunnel_url)
166
+
167
+ with self._lock:
168
+ self._loop = asyncio.get_running_loop()
169
+
170
+ current_delay = self.retry_delay
171
+
172
+ while not self._stop_event.is_set():
173
+ if self.api_key is None:
174
+ self._set_connected(False)
175
+ self._set_status("unauthenticated")
176
+ await self._wait_for_reconnect_signal()
177
+ continue
178
+
179
+ self._set_status("pending")
180
+ headers = self._build_connection_headers()
181
+
182
+ try:
183
+ logger.info("Connecting to %s", tunnel_url)
184
+ async with websockets.connect(
185
+ tunnel_url,
186
+ additional_headers=headers,
187
+ ssl=ssl_context,
188
+ proxy=None,
189
+ open_timeout=self.open_timeout,
190
+ ) as ws:
191
+ with self._lock:
192
+ self._active_ws = ws
193
+
194
+ current_delay = self.retry_delay
195
+ await self._consume_welcome(ws)
196
+
197
+ with self._lock:
198
+ self._reconnect_count = 0
199
+
200
+ send_lock = asyncio.Lock()
201
+ heartbeat_task = asyncio.create_task(
202
+ self._send_heartbeats(ws, send_lock)
203
+ )
204
+ try:
205
+ await self._handle_messages(ws, send_lock)
206
+ finally:
207
+ heartbeat_task.cancel()
208
+ await asyncio.gather(heartbeat_task, return_exceptions=True)
209
+ await self._close_all_websocket_sessions()
210
+ with self._lock:
211
+ self._active_ws = None
212
+ self._set_connected(False)
213
+ except websockets.exceptions.ConnectionClosed as exc:
214
+ self._set_connected(False)
215
+ self._set_status("failed")
216
+ logger.warning("Connection closed: %s", exc)
217
+ except Exception as exc: # pragma: no cover - network error path
218
+ self._set_connected(False)
219
+ self._set_status("failed")
220
+ logger.error("Connection error: %s", exc)
221
+ finally:
222
+ with self._lock:
223
+ self._active_ws = None
224
+
225
+ if self._stop_event.is_set():
226
+ break
227
+
228
+ if self._reconnect_event.is_set():
229
+ self._set_connected(False)
230
+ with self._lock:
231
+ self._reconnect_count = 0
232
+ self._reconnect_event.clear()
233
+ current_delay = self.retry_delay
234
+ self._set_status(self._status_for_auth())
235
+ logger.info("Manual reconnect requested. Reconnecting now.")
236
+ continue
237
+
238
+ with self._lock:
239
+ self._reconnect_count += 1
240
+ attempt = self._reconnect_count
241
+
242
+ if self.max_retries > 0 and attempt > self.max_retries:
243
+ self._set_status("failed")
244
+ logger.error(
245
+ "Max reconnection attempts (%s) reached. Giving up.",
246
+ self.max_retries,
247
+ )
248
+ break
249
+
250
+ logger.info("Reconnecting in %.1fs (attempt %s)", current_delay, attempt)
251
+
252
+ wait_time = 0.0
253
+ while (
254
+ wait_time < current_delay
255
+ and not self._stop_event.is_set()
256
+ and not self._reconnect_event.is_set()
257
+ ):
258
+ await asyncio.sleep(0.5)
259
+ wait_time += 0.5
260
+
261
+ if self._reconnect_event.is_set():
262
+ continue
263
+
264
+ current_delay = min(current_delay * 2, self._max_retry_delay)
265
+
266
+ with self._lock:
267
+ self._active_ws = None
268
+ self._loop = None
269
+ if self._status in {"connected", "pending"}:
270
+ self._status = self._status_for_auth()
271
+
272
+ logger.info("Disconnected from server.")
273
+
274
+ async def _consume_welcome(self, ws: Any) -> None:
275
+ """Read and process the initial welcome message when present."""
276
+ try:
277
+ welcome_message = await ws.recv()
278
+ welcome = json.loads(welcome_message)
279
+ except Exception:
280
+ self._set_connected(True)
281
+ return
282
+
283
+ if isinstance(welcome, dict) and welcome.get("type") == "worker_connected":
284
+ self._set_connected(
285
+ True,
286
+ worker_id=str(welcome.get("worker_id") or welcome.get("id") or ""),
287
+ worker_index=int(welcome.get("worker_index", -1)),
288
+ instance_id=str(welcome.get("instance_id") or "") or None,
289
+ tunnel_id=str(welcome.get("tunnel_id") or "") or None,
290
+ tunnel_name=str(welcome.get("tunnel_name") or "") or None,
291
+ max_workers=int(welcome.get("max_workers", -1)),
292
+ can_consume_tunnels=bool(
293
+ welcome.get("can_consume_tunnels", False)
294
+ ),
295
+ )
296
+ return
297
+
298
+ self._set_connected(True)
299
+
300
+ async def _handle_messages(
301
+ self, ws: Any, send_lock: asyncio.Lock
302
+ ) -> None:
303
+ """Handle request and websocket relay messages from the server."""
304
+ async with httpx.AsyncClient(timeout=Timeout(timeout=15.0)) as http_client:
305
+ while not self._stop_event.is_set() and not self._reconnect_event.is_set():
306
+ try:
307
+ message = await asyncio.wait_for(ws.recv(), timeout=1.0)
308
+ except asyncio.TimeoutError:
309
+ continue
310
+ except websockets.exceptions.ConnectionClosed:
311
+ break
312
+
313
+ try:
314
+ payload = json.loads(
315
+ message if isinstance(message, str) else message.decode("utf-8")
316
+ )
317
+ await self._handle_message(ws, http_client, payload, send_lock)
318
+ except Exception as exc:
319
+ logger.error("Error handling message: %s", exc)
320
+
321
+ async def _handle_message(
322
+ self,
323
+ server_ws: Any,
324
+ http_client: Any,
325
+ message: dict[str, Any],
326
+ send_lock: asyncio.Lock,
327
+ ) -> None:
328
+ """Dispatch one incoming server message to the matching handler."""
329
+ message_type = str(message.get("type", "")).strip().lower()
330
+ if message_type == "heartbeat_ack":
331
+ return
332
+
333
+ if message_type == "request":
334
+ response = await self._forward_request(http_client, message)
335
+ async with send_lock:
336
+ await server_ws.send(json.dumps(response))
337
+ return
338
+
339
+ if message_type == "ws_open":
340
+ await self._open_websocket_session(server_ws, message, send_lock)
341
+ return
342
+
343
+ if message_type == "ws_data":
344
+ await self._forward_websocket_data(message)
345
+ return
346
+
347
+ if message_type == "ws_close":
348
+ await self._close_websocket_session(message)
349
+
350
+ async def _forward_request(
351
+ self, client: Any, request: dict[str, Any]
352
+ ) -> dict[str, Any]:
353
+ """Forward one proxied HTTP request to the configured local target."""
354
+ request_id = str(request.get("id", ""))
355
+ method = str(request.get("method", "GET"))
356
+ path = str(request.get("path", "/"))
357
+ headers = request.get("headers", {})
358
+ body_b64 = str(request.get("body", ""))
359
+
360
+ body = base64.b64decode(body_b64) if body_b64 else None
361
+ url = urljoin(f"{self.target}/", path.lstrip("/"))
362
+ filtered_headers = _filter_hop_by_hop_headers(headers)
363
+
364
+ try:
365
+ response = await client.request(
366
+ method=method,
367
+ url=url,
368
+ headers=filtered_headers,
369
+ content=body,
370
+ follow_redirects=True,
371
+ )
372
+ return {
373
+ "type": "response",
374
+ "tunnel_id": str(request.get("tunnel_id") or self._tunnel_id or ""),
375
+ "id": request_id,
376
+ "status": response.status_code,
377
+ "headers": dict(response.headers.items()),
378
+ "body": base64.b64encode(response.content).decode("utf-8"),
379
+ "bodyEncoding": "base64",
380
+ }
381
+ except Exception as exc:
382
+ logger.error("Error forwarding request: %s", exc)
383
+ error_body = json.dumps({"error": str(exc)}).encode("utf-8")
384
+ return {
385
+ "type": "response",
386
+ "tunnel_id": str(request.get("tunnel_id") or self._tunnel_id or ""),
387
+ "id": request_id,
388
+ "status": 502,
389
+ "headers": {"Content-Type": "application/json"},
390
+ "body": base64.b64encode(error_body).decode("utf-8"),
391
+ "bodyEncoding": "base64",
392
+ "error": str(exc),
393
+ }
394
+
395
+ async def _open_websocket_session(
396
+ self,
397
+ server_ws: Any,
398
+ message: dict[str, Any],
399
+ send_lock: asyncio.Lock,
400
+ ) -> None:
401
+ """Open one websocket connection against the local target service."""
402
+ session_id = str(message.get("id", "")).strip()
403
+ if not session_id:
404
+ return
405
+
406
+ target_url = self._build_target_websocket_url(
407
+ path=str(message.get("path", "/")),
408
+ query_string=str(message.get("query_string", "") or ""),
409
+ )
410
+ headers = _filter_hop_by_hop_headers(message.get("headers", {}))
411
+
412
+ try:
413
+ local_ws = await websockets.connect(
414
+ target_url,
415
+ additional_headers=headers,
416
+ ssl=self._build_ssl_context(target_url),
417
+ proxy=None,
418
+ open_timeout=self.open_timeout,
419
+ )
420
+ except Exception as exc:
421
+ async with send_lock:
422
+ await server_ws.send(
423
+ json.dumps(
424
+ {
425
+ "type": "ws_error",
426
+ "tunnel_id": str(
427
+ message.get("tunnel_id") or self._tunnel_id or ""
428
+ ),
429
+ "id": session_id,
430
+ "error": str(exc),
431
+ }
432
+ )
433
+ )
434
+ return
435
+
436
+ self._ws_sessions[session_id] = local_ws
437
+ self._ws_reader_tasks[session_id] = asyncio.create_task(
438
+ self._relay_local_websocket(server_ws, session_id, local_ws, send_lock)
439
+ )
440
+ async with send_lock:
441
+ await server_ws.send(
442
+ json.dumps(
443
+ {
444
+ "type": "ws_opened",
445
+ "tunnel_id": str(
446
+ message.get("tunnel_id") or self._tunnel_id or ""
447
+ ),
448
+ "id": session_id,
449
+ }
450
+ )
451
+ )
452
+
453
+ async def _forward_websocket_data(self, message: dict[str, Any]) -> None:
454
+ """Forward one websocket frame from the server to the local target."""
455
+ session_id = str(message.get("id", "")).strip()
456
+ local_ws = self._ws_sessions.get(session_id)
457
+ if local_ws is None:
458
+ return
459
+
460
+ is_binary = bool(message.get("isBinary", False))
461
+ data = str(message.get("data", ""))
462
+ if is_binary:
463
+ await local_ws.send(base64.b64decode(data.encode("ascii")))
464
+ return
465
+
466
+ await local_ws.send(data)
467
+
468
+ async def _close_websocket_session(self, message: dict[str, Any]) -> None:
469
+ """Close one local websocket session at the server's request."""
470
+ session_id = str(message.get("id", "")).strip()
471
+ await self._cleanup_websocket_session(
472
+ session_id=session_id,
473
+ close_socket=True,
474
+ code=int(message.get("code") or 1000),
475
+ reason=str(message.get("reason") or ""),
476
+ )
477
+
478
+ async def _relay_local_websocket(
479
+ self,
480
+ server_ws: Any,
481
+ session_id: str,
482
+ local_ws: Any,
483
+ send_lock: asyncio.Lock,
484
+ ) -> None:
485
+ """Relay frames from one local websocket back to the server."""
486
+ try:
487
+ while True:
488
+ payload = await local_ws.recv()
489
+ if isinstance(payload, bytes):
490
+ message = {
491
+ "type": "ws_data",
492
+ "tunnel_id": self._tunnel_id or "",
493
+ "id": session_id,
494
+ "data": base64.b64encode(payload).decode("ascii"),
495
+ "isBinary": True,
496
+ }
497
+ else:
498
+ message = {
499
+ "type": "ws_data",
500
+ "tunnel_id": self._tunnel_id or "",
501
+ "id": session_id,
502
+ "data": payload,
503
+ "isBinary": False,
504
+ }
505
+ async with send_lock:
506
+ await server_ws.send(json.dumps(message))
507
+ except websockets.exceptions.ConnectionClosed as exc:
508
+ async with send_lock:
509
+ await server_ws.send(
510
+ json.dumps(
511
+ {
512
+ "type": "ws_close",
513
+ "tunnel_id": self._tunnel_id or "",
514
+ "id": session_id,
515
+ "code": exc.code,
516
+ "reason": exc.reason or None,
517
+ }
518
+ )
519
+ )
520
+ except Exception as exc:
521
+ async with send_lock:
522
+ await server_ws.send(
523
+ json.dumps(
524
+ {
525
+ "type": "ws_error",
526
+ "tunnel_id": self._tunnel_id or "",
527
+ "id": session_id,
528
+ "error": str(exc),
529
+ }
530
+ )
531
+ )
532
+ finally:
533
+ await self._cleanup_websocket_session(session_id=session_id)
534
+
535
+ async def _cleanup_websocket_session(
536
+ self,
537
+ session_id: str,
538
+ close_socket: bool = False,
539
+ code: int = 1000,
540
+ reason: str = "",
541
+ ) -> None:
542
+ """Remove one local websocket session and optionally close it."""
543
+ local_ws = self._ws_sessions.pop(session_id, None)
544
+ task = self._ws_reader_tasks.pop(session_id, None)
545
+
546
+ if close_socket and local_ws is not None:
547
+ try:
548
+ await local_ws.close(code=code, reason=reason)
549
+ except Exception:
550
+ pass
551
+
552
+ current_task = asyncio.current_task()
553
+ if task is not None and task is not current_task:
554
+ task.cancel()
555
+ await asyncio.gather(task, return_exceptions=True)
556
+
557
+ async def _close_all_websocket_sessions(self) -> None:
558
+ """Close all tracked local websocket sessions."""
559
+ for session_id in list(self._ws_sessions):
560
+ await self._cleanup_websocket_session(session_id, close_socket=True)
561
+
562
+ async def _send_heartbeats(self, server_ws: Any, send_lock: asyncio.Lock) -> None:
563
+ """Send periodic heartbeat messages while the worker socket is active."""
564
+ while not self._stop_event.is_set() and not self._reconnect_event.is_set():
565
+ await asyncio.sleep(HEARTBEAT_INTERVAL_SECONDS)
566
+ if self._tunnel_id is None or self._worker_id is None:
567
+ continue
568
+
569
+ heartbeat = {
570
+ "type": "heartbeat",
571
+ "tunnel_id": self._tunnel_id,
572
+ "id": self._worker_id,
573
+ "worker_id": self._worker_id,
574
+ }
575
+ try:
576
+ async with send_lock:
577
+ await server_ws.send(json.dumps(heartbeat))
578
+ except websockets.exceptions.ConnectionClosed:
579
+ return
580
+
581
+ def _build_tunnel_url(self) -> str:
582
+ """Build the worker websocket URL."""
583
+ return f"{self.host}/api/v1/tunnels/worker"
584
+
585
+ def _build_target_websocket_url(self, path: str, query_string: str = "") -> str:
586
+ """Build the local target websocket URL for one proxied session."""
587
+ parsed_target = urlparse(self.target)
588
+ scheme = parsed_target.scheme
589
+ if scheme == "http":
590
+ scheme = "ws"
591
+ elif scheme == "https":
592
+ scheme = "wss"
593
+
594
+ joined_url = urlparse(urljoin(f"{self.target}/", path.lstrip("/")))
595
+ return urlunparse(
596
+ (
597
+ scheme,
598
+ joined_url.netloc or parsed_target.netloc,
599
+ joined_url.path,
600
+ "",
601
+ query_string,
602
+ "",
603
+ )
604
+ )
605
+
606
+ def _set_connected(
607
+ self,
608
+ value: bool,
609
+ remote_address: Optional[str] = None,
610
+ remote_address_roundrobin: Optional[str] = None,
611
+ worker_index: int = -1,
612
+ pool_size: int = 0,
613
+ worker_id: Optional[str] = None,
614
+ instance_id: Optional[str] = None,
615
+ tunnel_id: Optional[str] = None,
616
+ tunnel_name: Optional[str] = None,
617
+ max_workers: int = -1,
618
+ can_consume_tunnels: bool = False,
619
+ ) -> None:
620
+ """Update connection state atomically."""
621
+ with self._lock:
622
+ self._connected = value
623
+ self._remote_address = remote_address if value else None
624
+ self._remote_address_roundrobin = (
625
+ remote_address_roundrobin if value else None
626
+ )
627
+ self._worker_index = worker_index if value else -1
628
+ self._pool_size = pool_size if value else 0
629
+ self._worker_id = worker_id if value else None
630
+ self._instance_id = instance_id if value else None
631
+ self._tunnel_id = tunnel_id if value else None
632
+ self._tunnel_name = tunnel_name if value else None
633
+ self._max_workers = max_workers if value else -1
634
+ self._can_consume_tunnels = can_consume_tunnels if value else False
635
+ if value:
636
+ self._status = "connected"
637
+ elif self.api_key is None:
638
+ self._status = "unauthenticated"
639
+
640
+ def _set_status(self, status: str) -> None:
641
+ """Store the current tunnel lifecycle status."""
642
+ with self._lock:
643
+ self._status = status
644
+
645
+ def _request_reconnect(self) -> None:
646
+ """Trigger a reconnect cycle and close the active websocket when needed."""
647
+ self._set_connected(False)
648
+ with self._lock:
649
+ self._reconnect_count = 0
650
+ ws = self._active_ws
651
+ loop = self._loop
652
+ self._status = self._status_for_auth()
653
+ self._reconnect_event.set()
654
+
655
+ if ws is not None and loop is not None and loop.is_running():
656
+ try:
657
+ future = asyncio.run_coroutine_threadsafe(
658
+ ws.close(code=1012, reason="Client reconnect"),
659
+ loop,
660
+ )
661
+ future.result(timeout=2)
662
+ except Exception as exc: # pragma: no cover - best-effort close path
663
+ logger.debug("Failed to close websocket during reconnect: %s", exc)
664
+
665
+ async def _wait_for_reconnect_signal(self) -> None:
666
+ """Wait until the client is reloaded, stopped, or resumed from an unauthenticated state."""
667
+ while not self._stop_event.is_set() and not self._reconnect_event.is_set():
668
+ await asyncio.sleep(0.5)
669
+
670
+ if self._reconnect_event.is_set():
671
+ self._reconnect_event.clear()
672
+
673
+ def _build_ssl_context(self, tunnel_url: str) -> Optional[ssl.SSLContext]:
674
+ """Create an SSL context for secure websocket URLs."""
675
+ if not tunnel_url.startswith("wss://"):
676
+ return None
677
+
678
+ ssl_context = ssl.create_default_context()
679
+ ssl_context.set_alpn_protocols(["http/1.1"])
680
+ if self.verify_ssl:
681
+ return ssl_context
682
+
683
+ ssl_context.check_hostname = False
684
+ ssl_context.verify_mode = ssl.CERT_NONE
685
+ return ssl_context
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "lange-python"
3
- version = "0.3.16"
3
+ version = "0.3.19"
4
4
  description = "A bundeld set of tools, clients for the lange-suite of tools and more."
5
5
  authors = [
6
6
  {name = "contact@robertlange.me"}
@@ -1,477 +0,0 @@
1
- from __future__ import annotations
2
-
3
- import asyncio
4
- import base64
5
- import json
6
- import logging
7
- import ssl
8
- import threading
9
- from typing import Any, Optional
10
- from urllib.parse import urljoin
11
-
12
- import httpx
13
- import websockets
14
- from httpx import Timeout
15
- from .._util import BaseLangeLabsClient
16
- from .._util._base_client import _UNSET_API_KEY
17
- from ._util import _filter_hop_by_hop_headers
18
-
19
- logger = logging.getLogger("lange.tunnel")
20
-
21
-
22
- class Tunnel(BaseLangeLabsClient):
23
- """
24
- Thread-based tunnel worker client.
25
-
26
- Connects to the tunnel runtime WebSocket endpoint and forwards incoming
27
- proxy messages to a local HTTP target.
28
-
29
- :param host: Base service URL, e.g. ``wss://example.com``.
30
- :param api_key: Bearer token used for worker authentication.
31
- :param target: Local HTTP target URL to forward tunnel traffic to.
32
- :param verify_ssl: Whether to verify TLS certificates for ``wss://`` hosts.
33
- :param max_retries: Maximum reconnect attempts, ``0`` for infinite retries.
34
- :param retry_delay: Initial reconnect delay in seconds.
35
- :param open_timeout: Timeout in seconds for the WebSocket opening handshake.
36
- :param daemon: Whether the worker thread is daemonized.
37
- :raises ValueError: If API key is provided but empty.
38
- """
39
-
40
- def __init__(
41
- self,
42
- host: str = "wss://tunnel.lange-labs.com",
43
- api_key: str | None = None,
44
- target: str = "http://localhost:80",
45
- verify_ssl: bool = True,
46
- max_retries: int = 5,
47
- retry_delay: float = 5.0,
48
- open_timeout: float = 20.0,
49
- daemon: bool = True,
50
- ) -> None:
51
- super().__init__(api_key=api_key, daemon=daemon, host=host)
52
- self.target = target.rstrip("/")
53
- self.verify_ssl = verify_ssl
54
- self.max_retries = max_retries
55
- self.retry_delay = retry_delay
56
- self.open_timeout = open_timeout
57
-
58
- self._max_retry_delay = 60.0
59
- self._stop_event = threading.Event()
60
- self._reconnect_event = threading.Event()
61
- self._connected = False
62
- self._remote_address: Optional[str] = None
63
- self._remote_address_roundrobin: Optional[str] = None
64
- self._worker_index = -1
65
- self._pool_size = 0
66
- self._reconnect_count = 0
67
- self._lock = threading.Lock()
68
- self._loop: Optional[asyncio.AbstractEventLoop] = None
69
- self._active_ws: Any = None
70
-
71
- @property
72
- def connected(self) -> bool:
73
- """
74
- Get current connection state.
75
-
76
- :returns: ``True`` when a worker socket is currently active.
77
- """
78
- with self._lock:
79
- return self._connected
80
-
81
- @property
82
- def status(self) -> str:
83
- """
84
- Get the current tunnel lifecycle status.
85
-
86
- :returns: Shared client status string.
87
- """
88
- with self._lock:
89
- return self._status
90
-
91
- @property
92
- def remote_address(self) -> Optional[str]:
93
- """
94
- Get worker-specific public address when provided by the server.
95
-
96
- :returns: Worker URL or ``None`` when unavailable.
97
- """
98
- with self._lock:
99
- return self._remote_address
100
-
101
- @property
102
- def remote_address_roundrobin(self) -> Optional[str]:
103
- """
104
- Get round-robin public address when provided by the server.
105
-
106
- :returns: Pool URL or ``None`` when unavailable.
107
- """
108
- with self._lock:
109
- return self._remote_address_roundrobin
110
-
111
- @property
112
- def worker_index(self) -> int:
113
- """
114
- Get current worker index from the latest welcome payload.
115
-
116
- :returns: Worker index, ``-1`` when disconnected.
117
- """
118
- with self._lock:
119
- return self._worker_index
120
-
121
- @property
122
- def pool_size(self) -> int:
123
- """
124
- Get current worker pool size from the latest welcome payload.
125
-
126
- :returns: Worker pool size, ``0`` when disconnected.
127
- """
128
- with self._lock:
129
- return self._pool_size
130
-
131
- @property
132
- def reconnect_count(self) -> int:
133
- """
134
- Get number of reconnect attempts after the last successful connection.
135
-
136
- :returns: Reconnect counter.
137
- """
138
- with self._lock:
139
- return self._reconnect_count
140
-
141
- def run(self) -> None:
142
- """
143
- Start the worker loop inside the thread.
144
-
145
- :returns: ``None``.
146
- """
147
- asyncio.run(self._run_async())
148
-
149
- def stop(self) -> None:
150
- """
151
- Request a graceful shutdown.
152
-
153
- :returns: ``None``.
154
- """
155
- self._stop_event.set()
156
-
157
- def reconnect(self) -> None:
158
- """
159
- Force a full reconnect cycle.
160
-
161
- :returns: ``None``.
162
- """
163
- self._request_reconnect()
164
-
165
- def reload(self, api_key: str | None | object = _UNSET_API_KEY) -> None:
166
- """
167
- Reload tunnel authentication and restart the connection flow when running.
168
-
169
- :param api_key: Optional replacement API key. Pass ``None`` explicitly to clear authentication.
170
- :returns: ``None``.
171
- """
172
- super().reload(api_key=api_key)
173
- self._set_connected(False)
174
-
175
- if self.is_alive() or self._loop is not None:
176
- self._request_reconnect()
177
- return
178
-
179
- self._set_status(self._status_for_auth())
180
-
181
- async def _run_async(self) -> None:
182
- """
183
- Run the worker connection loop with reconnect backoff.
184
-
185
- :returns: ``None``.
186
- """
187
- tunnel_url = self._build_tunnel_url()
188
- ssl_context = self._build_ssl_context(tunnel_url)
189
-
190
- with self._lock:
191
- self._loop = asyncio.get_running_loop()
192
-
193
- current_delay = self.retry_delay
194
-
195
- while not self._stop_event.is_set():
196
- if self.api_key is None:
197
- self._set_connected(False)
198
- self._set_status("unauthenticated")
199
- await self._wait_for_reconnect_signal()
200
- continue
201
-
202
- self._set_status("pending")
203
- headers = self._build_connection_headers()
204
-
205
- try:
206
- logger.info("Connecting to %s", tunnel_url)
207
- async with websockets.connect(
208
- tunnel_url,
209
- additional_headers=headers,
210
- ssl=ssl_context,
211
- proxy=None,
212
- open_timeout=self.open_timeout,
213
- ) as ws:
214
- with self._lock:
215
- self._active_ws = ws
216
-
217
- current_delay = self.retry_delay
218
- await self._consume_welcome(ws)
219
-
220
- with self._lock:
221
- self._reconnect_count = 0
222
-
223
- try:
224
- await self._handle_messages(ws)
225
- finally:
226
- with self._lock:
227
- self._active_ws = None
228
- self._set_connected(False)
229
- except websockets.exceptions.ConnectionClosed as exc:
230
- self._set_connected(False)
231
- self._set_status("failed")
232
- logger.warning("Connection closed: %s", exc)
233
- except Exception as exc: # pragma: no cover - network error path
234
- self._set_connected(False)
235
- self._set_status("failed")
236
- logger.error("Connection error: %s", exc)
237
- finally:
238
- with self._lock:
239
- self._active_ws = None
240
-
241
- if self._stop_event.is_set():
242
- break
243
-
244
- if self._reconnect_event.is_set():
245
- self._set_connected(False)
246
- with self._lock:
247
- self._reconnect_count = 0
248
- self._reconnect_event.clear()
249
- current_delay = self.retry_delay
250
- self._set_status(self._status_for_auth())
251
- logger.info("Manual reconnect requested. Reconnecting now.")
252
- continue
253
-
254
- with self._lock:
255
- self._reconnect_count += 1
256
- attempt = self._reconnect_count
257
-
258
- if self.max_retries > 0 and attempt > self.max_retries:
259
- self._set_status("failed")
260
- logger.error("Max reconnection attempts (%s) reached. Giving up.", self.max_retries)
261
- break
262
-
263
- logger.info("Reconnecting in %.1fs (attempt %s)", current_delay, attempt)
264
-
265
- wait_time = 0.0
266
- while (
267
- wait_time < current_delay
268
- and not self._stop_event.is_set()
269
- and not self._reconnect_event.is_set()
270
- ):
271
- await asyncio.sleep(0.5)
272
- wait_time += 0.5
273
-
274
- if self._reconnect_event.is_set():
275
- continue
276
-
277
- current_delay = min(current_delay * 2, self._max_retry_delay)
278
-
279
- with self._lock:
280
- self._active_ws = None
281
- self._loop = None
282
- if self._status in {"connected", "pending"}:
283
- self._status = self._status_for_auth()
284
-
285
- logger.info("Disconnected from server.")
286
-
287
- async def _consume_welcome(self, ws: Any) -> None:
288
- """
289
- Read and process the initial welcome message when present.
290
-
291
- :param ws: Active websocket client.
292
- :returns: ``None``.
293
- """
294
- try:
295
- welcome_message = await ws.recv()
296
- welcome = json.loads(welcome_message)
297
- except Exception:
298
- self._set_connected(True)
299
- return
300
-
301
- if isinstance(welcome, dict) and welcome.get("type") == "welcome":
302
- self._set_connected(
303
- True,
304
- remote_address=str(welcome.get("public_address") or "") or None,
305
- remote_address_roundrobin=str(welcome.get("public_address_generic") or "") or None,
306
- worker_index=int(welcome.get("worker_index", -1)),
307
- pool_size=int(welcome.get("pool_size", 0)),
308
- )
309
- return
310
-
311
- self._set_connected(True)
312
-
313
- async def _handle_messages(self, ws: Any) -> None:
314
- """
315
- Handle request/response proxy messages over the worker websocket.
316
-
317
- :param ws: Active websocket client.
318
- :returns: ``None``.
319
- """
320
- async with httpx.AsyncClient(timeout=Timeout(timeout=15.0)) as http_client:
321
- while not self._stop_event.is_set() and not self._reconnect_event.is_set():
322
- try:
323
- message = await asyncio.wait_for(ws.recv(), timeout=1.0)
324
- request = json.loads(message)
325
- response = await self._forward_request(http_client, request)
326
- await ws.send(json.dumps(response))
327
- except asyncio.TimeoutError:
328
- continue
329
- except websockets.exceptions.ConnectionClosed:
330
- break
331
- except Exception as exc:
332
- logger.error("Error handling message: %s", exc)
333
-
334
- async def _forward_request(self, client: Any, request: dict[str, Any]) -> dict[str, Any]:
335
- """
336
- Forward a single proxied request to the configured local target.
337
-
338
- :param client: HTTP client supporting ``request``.
339
- :param request: Incoming tunnel request payload.
340
- :returns: Tunnel response payload.
341
- """
342
- request_id = str(request.get("id", ""))
343
- method = str(request.get("method", "GET"))
344
- path = str(request.get("path", "/"))
345
- headers = request.get("headers", {})
346
- body_b64 = str(request.get("body", ""))
347
-
348
- body = base64.b64decode(body_b64) if body_b64 else None
349
-
350
- url = urljoin(f"{self.target}/", path.lstrip("/"))
351
- filtered_headers = _filter_hop_by_hop_headers(headers)
352
-
353
- try:
354
- response = await client.request(
355
- method=method,
356
- url=url,
357
- headers=filtered_headers,
358
- content=body,
359
- follow_redirects=True,
360
- )
361
- return {
362
- "id": request_id,
363
- "status": response.status_code,
364
- "headers": dict(response.headers.items()),
365
- "body": base64.b64encode(response.content).decode("utf-8"),
366
- }
367
- except Exception as exc:
368
- logger.error("Error forwarding request: %s", exc)
369
- error_body = json.dumps({"error": str(exc)}).encode("utf-8")
370
- return {
371
- "id": request_id,
372
- "status": 502,
373
- "headers": {"Content-Type": "application/json"},
374
- "body": base64.b64encode(error_body).decode("utf-8"),
375
- }
376
-
377
- def _build_tunnel_url(self) -> str:
378
- """
379
- Build the worker websocket URL.
380
-
381
- :returns: WebSocket endpoint URL.
382
- """
383
- return f"{self.host}/api/tunnels/connection"
384
-
385
- def _set_connected(
386
- self,
387
- value: bool,
388
- remote_address: Optional[str] = None,
389
- remote_address_roundrobin: Optional[str] = None,
390
- worker_index: int = -1,
391
- pool_size: int = 0,
392
- ) -> None:
393
- """
394
- Update connection state atomically.
395
-
396
- :param value: New connected state.
397
- :param remote_address: Worker URL from welcome message.
398
- :param remote_address_roundrobin: Pool URL from welcome message.
399
- :param worker_index: Worker index from welcome message.
400
- :param pool_size: Pool size from welcome message.
401
- :returns: ``None``.
402
- """
403
- with self._lock:
404
- self._connected = value
405
- self._remote_address = remote_address if value else None
406
- self._remote_address_roundrobin = remote_address_roundrobin if value else None
407
- self._worker_index = worker_index if value else -1
408
- self._pool_size = pool_size if value else 0
409
- if value:
410
- self._status = "connected"
411
- elif self.api_key is None:
412
- self._status = "unauthenticated"
413
-
414
- def _set_status(self, status: str) -> None:
415
- """
416
- Store the current tunnel lifecycle status.
417
-
418
- :param status: New lifecycle status.
419
- :returns: ``None``.
420
- """
421
- with self._lock:
422
- self._status = status
423
-
424
- def _request_reconnect(self) -> None:
425
- """
426
- Trigger a reconnect cycle and close the active websocket when needed.
427
-
428
- :returns: ``None``.
429
- """
430
- self._set_connected(False)
431
- with self._lock:
432
- self._reconnect_count = 0
433
- ws = self._active_ws
434
- loop = self._loop
435
- self._status = self._status_for_auth()
436
- self._reconnect_event.set()
437
-
438
- if ws is not None and loop is not None and loop.is_running():
439
- try:
440
- future = asyncio.run_coroutine_threadsafe(
441
- ws.close(code=1012, reason="Client reconnect"),
442
- loop,
443
- )
444
- future.result(timeout=2)
445
- except Exception as exc: # pragma: no cover - best-effort close path
446
- logger.debug("Failed to close websocket during reconnect: %s", exc)
447
-
448
- async def _wait_for_reconnect_signal(self) -> None:
449
- """
450
- Wait until the client is reloaded, stopped, or resumed from an unauthenticated state.
451
-
452
- :returns: ``None``.
453
- """
454
- while not self._stop_event.is_set() and not self._reconnect_event.is_set():
455
- await asyncio.sleep(0.5)
456
-
457
- if self._reconnect_event.is_set():
458
- self._reconnect_event.clear()
459
-
460
- def _build_ssl_context(self, tunnel_url: str) -> Optional[ssl.SSLContext]:
461
- """
462
- Create SSL context for secure websocket URLs.
463
-
464
- :param tunnel_url: Resolved websocket URL.
465
- :returns: SSL context for ``wss`` or ``None`` for plain ``ws``.
466
- """
467
- if not tunnel_url.startswith("wss://"):
468
- return None
469
-
470
- ssl_context = ssl.create_default_context()
471
- ssl_context.set_alpn_protocols(["http/1.1"])
472
- if self.verify_ssl:
473
- return ssl_context
474
-
475
- ssl_context.check_hostname = False
476
- ssl_context.verify_mode = ssl.CERT_NONE
477
- return ssl_context