130 lines
4.7 KiB
Python
130 lines
4.7 KiB
Python
"""
|
|
PiPedal WebSocket proxy.
|
|
|
|
Connects to PiPedal's WebSocket at ws://<host>:8080/pipedal and relays
|
|
messages to/from the browser via FastAPI's WebSocket endpoint.
|
|
|
|
Protocol (PiPedal WS):
|
|
Messages are JSON arrays: [{"message": "msgName", "replyTo": N}, optionalBody]
|
|
Replies: [{"reply": N, "message": "msgName"}, optionalBody]
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from typing import Optional
|
|
|
|
import websockets
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class PiPedalWSProxy:
|
|
"""Manages a single upstream WebSocket connection to PiPedal and relays
|
|
messages to a single downstream (browser) WebSocket."""
|
|
|
|
def __init__(self, host: str = "192.168.0.245", port: int = 8080, path: str = "/pipedal"):
|
|
self._host = host
|
|
self._port = port
|
|
self._path = path
|
|
self._upstream_ws: Optional[websockets.WebSocketClientProtocol] = None
|
|
self._downstream_ws: Optional["WebSocket"] = None # fastapi WebSocket
|
|
self._task: Optional[asyncio.Task] = None
|
|
self._connected = False
|
|
self._lock = asyncio.Lock()
|
|
|
|
@property
|
|
def upstream_url(self) -> str:
|
|
return f"ws://{self._host}:{self._port}{self._path}"
|
|
|
|
@property
|
|
def is_connected(self) -> bool:
|
|
return self._connected
|
|
|
|
async def connect_upstream(self) -> bool:
|
|
"""Connect to PiPedal's WebSocket. Returns True on success."""
|
|
try:
|
|
self._upstream_ws = await websockets.connect(
|
|
self.upstream_url,
|
|
ping_interval=30,
|
|
ping_timeout=10,
|
|
close_timeout=5,
|
|
)
|
|
self._connected = True
|
|
logger.info("Connected to PiPedal WS at %s", self.upstream_url)
|
|
return True
|
|
except Exception as e:
|
|
logger.warning("Failed to connect to PiPedal WS: %s", e)
|
|
self._connected = False
|
|
return False
|
|
|
|
async def close_upstream(self):
|
|
"""Close upstream connection."""
|
|
self._connected = False
|
|
if self._upstream_ws:
|
|
try:
|
|
await self._upstream_ws.close()
|
|
except Exception:
|
|
pass
|
|
self._upstream_ws = None
|
|
|
|
async def relay_loop(self, downstream_ws):
|
|
"""Run relay loop: forward upstream→downstream and downstream→upstream."""
|
|
self._downstream_ws = downstream_ws
|
|
|
|
async def upstream_to_downstream():
|
|
"""Read from PiPedal WS and forward to browser."""
|
|
try:
|
|
async for message in self._upstream_ws:
|
|
try:
|
|
await self._downstream_ws.send_text(message)
|
|
except Exception:
|
|
logger.warning("Downstream send failed, stopping relay")
|
|
break
|
|
except websockets.exceptions.ConnectionClosed:
|
|
logger.info("Upstream connection closed")
|
|
except Exception as e:
|
|
logger.warning("Upstream→downstream relay error: %s", e)
|
|
finally:
|
|
self._connected = False
|
|
|
|
async def downstream_to_upstream():
|
|
"""Read from browser and forward to PiPedal WS."""
|
|
try:
|
|
while True:
|
|
raw = await self._downstream_ws.receive_text()
|
|
if self._upstream_ws and self._connected:
|
|
await self._upstream_ws.send(raw)
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
self._connected = False
|
|
|
|
# Run both directions concurrently
|
|
await asyncio.gather(
|
|
upstream_to_downstream(),
|
|
downstream_to_upstream(),
|
|
)
|
|
|
|
async def send_to_pipedal(self, message: str):
|
|
"""Send a raw message to PiPedal upstream."""
|
|
async with self._lock:
|
|
if self._upstream_ws and self._connected:
|
|
await self._upstream_ws.send(message)
|
|
|
|
async def request(self, message_name: str, body=None) -> dict:
|
|
"""Send a request to PiPedal and wait for reply using replyTo/reply mechanism."""
|
|
reply_to = id(message_name) & 0xFFFFFFFF # simple unique id
|
|
msg = [{"message": message_name, "replyTo": reply_to}]
|
|
if body is not None:
|
|
msg.append(body)
|
|
|
|
future: asyncio.Future = asyncio.get_running_loop().create_future()
|
|
# We need a reply handler - but since this is a simple proxy,
|
|
# we use the REST endpoint approach instead for discovery.
|
|
# For inline use, we'd need pending_request tracking.
|
|
await self.send_to_pipedal(json.dumps(msg))
|
|
# Note: real reply matching requires tracking in the relay loop.
|
|
# This is a simplified version.
|
|
return {"status": "sent", "reply_to": reply_to}
|