borecli 1.0.0__py3-none-any.whl

bore/tunnel/heartbeat.py

@@ -0,0 +1,154 @@
+# borecli/bore/tunnel/heartbeat.py
+
+import asyncio
+import json
+import logging
+import time
+from typing import Optional
+
+from bore.frames import make_frame
+from bore.protocol import MessageType
+
+logger = logging.getLogger(__name__)
+
+
+class Heartbeat:
+    """
+    Maintains the BoreHook tunnel heartbeat.
+
+    Responsibilities
+    ----------------
+    - Send periodic ping frames
+    - Record pong replies
+    - Detect dead connections
+    - Allow the tunnel client to reconnect
+    """
+
+    def __init__(
+        self,
+        websocket,
+        *,
+        interval: int = 15,
+        timeout: int = 45,
+    ):
+        self.websocket = websocket
+
+        self.interval = interval
+        self.timeout = timeout
+
+        self._running = False
+        self._task: Optional[asyncio.Task] = None
+
+        self._last_pong = time.monotonic()
+
+    async def start(self):
+        """
+        Start the heartbeat loop.
+        """
+
+        if self._running:
+            return
+
+        logger.info("Heartbeat started")
+
+        self._running = True
+
+        self._task = asyncio.create_task(
+            self._run(),
+            name="heartbeat",
+        )
+
+    async def stop(self):
+        """
+        Stop the heartbeat loop.
+        """
+
+        self._running = False
+
+        if self._task:
+
+            self._task.cancel()
+
+            try:
+                await self._task
+
+            except asyncio.CancelledError:
+                pass
+
+        logger.info("Heartbeat stopped")
+
+    def pong(self):
+        """
+        Called by the dispatcher whenever
+        a pong frame is received.
+        """
+
+        self._last_pong = time.monotonic()
+
+        logger.debug("Heartbeat pong")
+
+    async def _run(self):
+        """
+        Internal heartbeat loop.
+        """
+
+        while self._running:
+
+            try:
+
+                #
+                # Send ping
+                #
+
+                frame = make_frame(
+                    MessageType.PING,
+                )
+
+                await self.websocket.send(
+                    json.dumps(frame),
+                )
+
+                logger.debug("Heartbeat ping")
+
+                #
+                # Has the server stopped responding?
+                #
+
+                elapsed = (
+                    time.monotonic()
+                    - self._last_pong
+                )
+
+                if elapsed > self.timeout:
+
+                    logger.warning(
+                        "Heartbeat timeout "
+                        "(%.1fs)",
+                        elapsed,
+                    )
+
+                    await self.websocket.close()
+
+                    break
+
+                await asyncio.sleep(
+                    self.interval,
+                )
+
+            except asyncio.CancelledError:
+                break
+
+            except Exception as exc:
+
+                logger.exception(
+                    "Heartbeat error: %s",
+                    exc,
+                )
+
+                try:
+                    await self.websocket.close()
+
+                except Exception:
+                    pass
+
+                break

bore/tunnel.py

@@ -0,0 +1,43 @@
+import requests
+
+from .config import get_token
+
+
+# API_URL = "http://127.0.0.1:8000"
+API_URL = "https://api.borehook.com"
+
+
+def start_tunnel(subdomain, port):
+
+    token = get_token()
+    print("TOKEN:", token)
+
+    if not token:
+        raise Exception(
+            "Run: bore login"
+        )
+
+    response = requests.post(
+        f"{API_URL}/api/tunnels/start/",
+        json={
+            "subdomain": subdomain,
+            "port": port,
+        },
+        # headers={
+        #     "Authorization": f"Bearer {token}"
+        # },
+        headers = {
+              "Authorization": f"Token {token}"
+   },
+        timeout=30,
+    )
+    
+    # print("STATUS:", response.status_code)
+    # print("BODY:", response.text)
+
+    response.raise_for_status()
+
+    data = response.json()
+    print(response.json())
+
+    return data

bore/websocket/connection.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+import asyncio
+import time
+import uuid
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Optional
+
+from websockets.client import WebSocketClientProtocol
+
+
+class ConnectionState(str, Enum):
+    CONNECTING = "connecting"
+    CONNECTED = "connected"
+    CLOSING = "closing"
+    CLOSED = "closed"
+
+
+@dataclass(slots=True)
+class ConnectionStats:
+    bytes_sent: int = 0
+    bytes_received: int = 0
+
+    messages_sent: int = 0
+    messages_received: int = 0
+
+    created_at: float = field(default_factory=time.time)
+    last_activity: float = field(default_factory=time.time)
+
+    def sent(self, message):
+        size = len(message) if isinstance(message, bytes) else len(str(message))
+
+        self.bytes_sent += size
+        self.messages_sent += 1
+        self.last_activity = time.time()
+
+    def received(self, message):
+        size = len(message) if isinstance(message, bytes) else len(str(message))
+
+        self.bytes_received += size
+        self.messages_received += 1
+        self.last_activity = time.time()
+
+    @property
+    def uptime(self):
+        return time.time() - self.created_at
+
+    @property
+    def idle(self):
+        return time.time() - self.last_activity
+
+
+@dataclass(slots=True)
+class WebSocketConnection:
+    """
+    Represents one localhost websocket connection.
+
+    This class stores metadata only.
+    """
+
+    websocket: WebSocketClientProtocol
+    url: str
+
+    connection_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+
+    tunnel_id: Optional[str] = None
+
+    state: ConnectionState = ConnectionState.CONNECTING
+
+    stats: ConnectionStats = field(default_factory=ConnectionStats)
+
+    send_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+
+    metadata: dict = field(default_factory=dict)
+
+    @property
+    def connected(self) -> bool:
+        return self.state == ConnectionState.CONNECTED
+
+    def mark_connected(self):
+        self.state = ConnectionState.CONNECTED
+
+    def mark_closing(self):
+        self.state = ConnectionState.CLOSING
+
+    def mark_closed(self):
+        self.state = ConnectionState.CLOSED
+
+    def record_sent(self, message):
+        self.stats.sent(message)
+
+    def record_received(self, message):
+        self.stats.received(message)
+
+    def to_dict(self):
+        return {
+            "connection_id": self.connection_id,
+            "tunnel_id": self.tunnel_id,
+            "url": self.url,
+            "state": self.state.value,
+            "uptime": self.stats.uptime,
+            "idle": self.stats.idle,
+            "messages_sent": self.stats.messages_sent,
+            "messages_received": self.stats.messages_received,
+            "bytes_sent": self.stats.bytes_sent,
+            "bytes_received": self.stats.bytes_received,
+        }

bore/websocket/frames.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+import base64
+import json
+import time
+import uuid
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Optional
+
+
+class FrameType(str, Enum):
+    TEXT = "text"
+    BINARY = "binary"
+    PING = "ping"
+    PONG = "pong"
+    CLOSE = "close"
+
+
+class Direction(str, Enum):
+    CLIENT_TO_LOCAL = "client_to_local"
+    LOCAL_TO_CLIENT = "local_to_client"
+
+
+@dataclass(slots=True)
+class WebSocketFrame:
+    """
+    Represents a websocket frame flowing through the proxy.
+
+    This class is transport-independent and can be used for
+    logging, inspection, replay, and metrics.
+    """
+
+    payload: str | bytes
+
+    frame_type: FrameType
+
+    direction: Direction
+
+    connection_id: Optional[str] = None
+
+    timestamp: float = field(default_factory=time.time)
+
+    frame_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+
+    @property
+    def is_text(self) -> bool:
+        return self.frame_type == FrameType.TEXT
+
+    @property
+    def is_binary(self) -> bool:
+        return self.frame_type == FrameType.BINARY
+
+    @property
+    def size(self) -> int:
+        if isinstance(self.payload, bytes):
+            return len(self.payload)
+
+        return len(self.payload.encode())
+
+    def as_json(self) -> Optional[Any]:
+        """
+        Return parsed JSON if payload contains JSON.
+        """
+
+        if not self.is_text:
+            return None
+
+        try:
+            return json.loads(self.payload)
+        except Exception:
+            return None
+
+    def serialize(self) -> dict:
+        """
+        Convert to a JSON-safe dictionary.
+        """
+
+        if self.is_binary:
+            payload = base64.b64encode(self.payload).decode()
+            encoding = "base64"
+        else:
+            payload = self.payload
+            encoding = "utf-8"
+
+        return {
+            "frame_id": self.frame_id,
+            "connection_id": self.connection_id,
+            "timestamp": self.timestamp,
+            "direction": self.direction.value,
+            "type": self.frame_type.value,
+            "size": self.size,
+            "encoding": encoding,
+            "payload": payload,
+        }
+
+    @classmethod
+    def text(
+        cls,
+        payload: str,
+        direction: Direction,
+        connection_id: str | None = None,
+    ):
+        return cls(
+            payload=payload,
+            frame_type=FrameType.TEXT,
+            direction=direction,
+            connection_id=connection_id,
+        )
+
+    @classmethod
+    def binary(
+        cls,
+        payload: bytes,
+        direction: Direction,
+        connection_id: str | None = None,
+    ):
+        return cls(
+            payload=payload,
+            frame_type=FrameType.BINARY,
+            direction=direction,
+            connection_id=connection_id,
+        )
+
+    @classmethod
+    def ping(
+        cls,
+        payload: bytes = b"",
+        connection_id: str | None = None,
+    ):
+        return cls(
+            payload=payload,
+            frame_type=FrameType.PING,
+            direction=Direction.CLIENT_TO_LOCAL,
+            connection_id=connection_id,
+        )
+
+    @classmethod
+    def pong(
+        cls,
+        payload: bytes = b"",
+        connection_id: str | None = None,
+    ):
+        return cls(
+            payload=payload,
+            frame_type=FrameType.PONG,
+            direction=Direction.LOCAL_TO_CLIENT,
+            connection_id=connection_id,
+        )
+
+    @classmethod
+    def close(
+        cls,
+        payload: str = "",
+        connection_id: str | None = None,
+    ):
+        return cls(
+            payload=payload,
+            frame_type=FrameType.CLOSE,
+            direction=Direction.CLIENT_TO_LOCAL,
+            connection_id=connection_id,
+        )

bore/websocket/manager.py

@@ -0,0 +1,174 @@
+import asyncio
+import logging
+import time
+from dataclasses import dataclass, field
+from typing import Dict, Optional
+
+import websockets
+from websockets.client import WebSocketClientProtocol
+from websockets.exceptions import ConnectionClosed
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ManagedConnection:
+    connection_id: str
+    websocket: WebSocketClientProtocol
+    url: str
+
+    created_at: float = field(default_factory=time.time)
+    last_activity: float = field(default_factory=time.time)
+
+    bytes_sent: int = 0
+    bytes_received: int = 0
+
+    send_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
+
+    @property
+    def connected(self) -> bool:
+        return not self.websocket.closed
+
+
+class WebSocketManager:
+    """
+    Manages websocket connections from the CLI to localhost.
+
+    This class DOES NOT proxy websocket traffic.
+    It only manages the lifecycle of localhost websocket connections.
+    """
+
+    def __init__(self):
+        self._connections: Dict[str, ManagedConnection] = {}
+        self._lock = asyncio.Lock()
+
+    async def connect(
+        self,
+        connection_id: str,
+        local_port: int,
+        path: str,
+        query: str = "",
+        headers: Optional[dict] = None,
+    ) -> ManagedConnection:
+
+        url = f"ws://127.0.0.1:{local_port}{path}"
+
+        if query:
+            url += f"?{query}"
+
+        logger.info("Connecting to %s", url)
+
+        ws = await websockets.connect(
+            url,
+            additional_headers=headers or {},
+            ping_interval=20,
+            ping_timeout=20,
+            close_timeout=5,
+            max_size=None,
+            max_queue=1024,
+        )
+
+        connection = ManagedConnection(
+            connection_id=connection_id,
+            websocket=ws,
+            url=url,
+        )
+
+        async with self._lock:
+            self._connections[connection_id] = connection
+
+        logger.info("Connected websocket %s", connection_id)
+
+        return connection
+
+    async def send(
+        self,
+        connection_id: str,
+        message,
+    ):
+
+        conn = self.get(connection_id)
+
+        if conn is None:
+            raise KeyError(connection_id)
+
+        async with conn.send_lock:
+
+            await conn.websocket.send(message)
+
+            size = len(message) if isinstance(message, bytes) else len(str(message))
+
+            conn.bytes_sent += size
+            conn.last_activity = time.time()
+
+    async def recv(
+        self,
+        connection_id: str,
+    ):
+
+        conn = self.get(connection_id)
+
+        if conn is None:
+            raise KeyError(connection_id)
+
+        try:
+
+            message = await conn.websocket.recv()
+
+            size = len(message) if isinstance(message, bytes) else len(str(message))
+
+            conn.bytes_received += size
+            conn.last_activity = time.time()
+
+            return message
+
+        except ConnectionClosed:
+            await self.close(connection_id)
+            raise
+
+    async def close(self, connection_id: str):
+
+        async with self._lock:
+            conn = self._connections.pop(connection_id, None)
+
+        if conn is None:
+            return
+
+        try:
+            await conn.websocket.close()
+        except Exception:
+            logger.exception("Failed closing websocket %s", connection_id)
+
+    async def close_all(self):
+
+        ids = list(self._connections.keys())
+
+        await asyncio.gather(
+            *(self.close(cid) for cid in ids),
+            return_exceptions=True,
+        )
+
+    def get(self, connection_id: str) -> Optional[ManagedConnection]:
+        return self._connections.get(connection_id)
+
+    def exists(self, connection_id: str) -> bool:
+        return connection_id in self._connections
+
+    def count(self) -> int:
+        return len(self._connections)
+
+    def stats(self):
+
+        now = time.time()
+
+        return {
+            cid: {
+                "url": conn.url,
+                "connected": conn.connected,
+                "uptime": now - conn.created_at,
+                "idle": now - conn.last_activity,
+                "bytes_sent": conn.bytes_sent,
+                "bytes_received": conn.bytes_received,
+            }
+            for cid, conn in self._connections.items()
+        }