Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
129 changes: 93 additions & 36 deletions dataplicity_cli/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -605,6 +605,31 @@ def _friendly_response_message(default_message: str, response_data: Any, respons
return message


def _debug_api_request(
state: AppContext,
*,
method: str,
path: str,
params: Optional[Dict[str, Any]] = None,
json_data: Optional[Dict[str, Any]] = None,
data: Optional[Dict[str, Any]] = None,
) -> None:
if not state.debug:
return
if path.startswith("http://") or path.startswith("https://"):
url = path
else:
url = f"{state.config.base_url.rstrip('/')}/{path.lstrip('/')}"
payload: Dict[str, Any] = {"method": method.upper(), "url": url}
if params:
payload["params"] = _sanitize_payload(params)
if json_data:
payload["json"] = _sanitize_payload(json_data)
if data:
payload["data"] = _sanitize_payload(data)
sys.stderr.write(f"[debug] API request {json.dumps(payload, sort_keys=True)}\n")


def _incident_automation_org_hash(endpoint: str) -> str:
match = re.search(r"/api/organisations/([^/]+)/", endpoint)
if not match:
Expand Down Expand Up @@ -919,7 +944,7 @@ class _PortForwardDashboard:
active_clients: int = 0
total_connections: int = 0
rejected_connections: int = 0
current_connection_started_at: Optional[float] = None
connection_started_at: Dict[str, float] = field(default_factory=dict)
first_byte_latencies_ms: List[float] = field(default_factory=list)
protocols_seen: List[str] = field(default_factory=list)
listener_ready: bool = False
Expand All @@ -940,19 +965,25 @@ def apply_event(self, event: Any) -> None:
self.stop_reason = "stopped"
return
if kind == "connection_opened":
self.active_clients = 1
self.active_clients += 1
self.total_connections += 1
self.current_connection_started_at = timestamp
self.connection_started_at[detail] = timestamp
return
if kind == "connection_closed":
self.active_clients = 0
self.current_connection_started_at = None
self.active_clients = max(0, self.active_clients - 1)
if detail:
self.connection_started_at.pop(detail, None)
return
if kind == "connection_rejected":
self.rejected_connections += 1
return
if kind == "first_remote_byte" and self.current_connection_started_at is not None:
latency_ms = (timestamp - self.current_connection_started_at) * 1000.0
if kind == "first_remote_byte":
started_at = self.connection_started_at.get(detail)
if started_at is None and self.connection_started_at:
started_at = next(iter(self.connection_started_at.values()))
if started_at is None:
return
latency_ms = (timestamp - started_at) * 1000.0
self.first_byte_latencies_ms.append(latency_ms)
self.first_byte_latencies_ms = self.first_byte_latencies_ms[-64:]
return
Expand Down Expand Up @@ -2206,7 +2237,9 @@ def devices_history_list(
if include_deleted:
params["include_deleted"] = "1"

response = state.api.get(f"/api/developer/devices/{resolved_hash}/timeline/", params=params)
endpoint = f"/api/developer/devices/{resolved_hash}/timeline/"
_debug_api_request(state, method="GET", path=endpoint, params=params)
response = state.api.get(endpoint, params=params)
if not response.ok:
message = _friendly_response_message("Unable to load device history.", response.data, response.text)
if state.json_output:
Expand Down Expand Up @@ -2270,7 +2303,9 @@ def devices_history_show(
params: Dict[str, Any] = {"limit": 500}
if include_deleted:
params["include_deleted"] = "1"
response = state.api.get(f"/api/developer/devices/{resolved_hash}/timeline/", params=params)
endpoint = f"/api/developer/devices/{resolved_hash}/timeline/"
_debug_api_request(state, method="GET", path=endpoint, params=params)
response = state.api.get(endpoint, params=params)
if not response.ok:
message = _friendly_response_message("Unable to load device history.", response.data, response.text)
if state.json_output:
Expand Down Expand Up @@ -2361,7 +2396,9 @@ def devices_history_comment(
_show_error(state.console, message)
raise typer.Exit(code=2)
payload["type"] = normalized_type
response = state.api.post(f"/api/developer/devices/{resolved_hash}/timeline/", json_data=payload)
endpoint = f"/api/developer/devices/{resolved_hash}/timeline/"
_debug_api_request(state, method="POST", path=endpoint, json_data=payload)
response = state.api.post(endpoint, json_data=payload)
if not response.ok:
message = _friendly_response_message("Unable to add device history comment.", response.data, response.text)
if state.json_output:
Expand Down Expand Up @@ -2399,9 +2436,11 @@ def devices_history_delete(
state = _ctx(ctx)
_require_auth(state)
resolved_hash = _resolve_device_hash_interactive(state, device_hash, action_name="history delete")
endpoint = f"/api/developer/devices/{resolved_hash}/timeline/{event_id}/"
_debug_api_request(state, method="DELETE", path=endpoint)
response = state.api.request(
"DELETE",
f"/api/developer/devices/{resolved_hash}/timeline/{event_id}/",
endpoint,
)
if not response.ok:
message = _friendly_response_message("Unable to mark history entry as deleted.", response.data, response.text)
Expand Down Expand Up @@ -2784,24 +2823,32 @@ async def runner() -> Dict[str, Any]:
await m2m.connect()
try:
identity = await m2m.wait_for_identity()
response = _open_remote_access_port(
state,
resolved_hash,
{
"m2m_identity": identity,
"service": "redirect-port",
"port": remote_port,
},
)
if not response.ok:
detail = _friendly_response_message("Unable to open port redirect", response.data, response.text)
raise RuntimeError(detail)
channel_port = await m2m.wait_for_channel_open()
open_channel_lock = asyncio.Lock()

async def open_redirect_channel() -> int:
async with open_channel_lock:
response = await asyncio.to_thread(
_open_remote_access_port,
state,
resolved_hash,
{
"m2m_identity": identity,
"service": "redirect-port",
"port": remote_port,
},
)
if not response.ok:
detail = _friendly_response_message("Unable to open port redirect", response.data, response.text)
raise RuntimeError(detail)
return await m2m.wait_for_channel_open()

channel_port = await open_redirect_channel()
forward_task = asyncio.create_task(
run_port_forward(
m2m,
channel_port,
local_port,
channel_factory=open_redirect_channel,
event_callback=dashboard.apply_event,
)
)
Expand Down Expand Up @@ -2912,25 +2959,35 @@ def on_forward_event(event: Any) -> None:

try:
identity = await m2m.wait_for_identity(timeout=float(connect_timeout))
response = _open_remote_access_port(
state,
resolved_hash,
{
"m2m_identity": identity,
"service": "redirect-port",
"port": remote_port,
},
)
if not response.ok:
raise RuntimeError(_friendly_response_message("Unable to open SSH tunnel.", response.data, response.text))
channel_port = await m2m.wait_for_channel_open(timeout=float(connect_timeout))
open_channel_lock = asyncio.Lock()

async def open_redirect_channel() -> int:
async with open_channel_lock:
response = await asyncio.to_thread(
_open_remote_access_port,
state,
resolved_hash,
{
"m2m_identity": identity,
"service": "redirect-port",
"port": remote_port,
},
)
if not response.ok:
raise RuntimeError(
_friendly_response_message("Unable to open SSH tunnel.", response.data, response.text)
)
return await m2m.wait_for_channel_open(timeout=float(connect_timeout))

channel_port = await open_redirect_channel()
from .remote_access import run_port_forward

forward_task = asyncio.create_task(
run_port_forward(
m2m,
channel_port,
resolved_local_port,
channel_factory=open_redirect_channel,
event_callback=on_forward_event,
)
)
Expand Down
65 changes: 46 additions & 19 deletions dataplicity_cli/remote_access.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
import time
import tty
from dataclasses import dataclass
from typing import Callable, Optional
from typing import Awaitable, Callable, Optional

from .m2m import M2MClient

Expand Down Expand Up @@ -43,6 +43,7 @@ class PortForwardEvent:


PortForwardEventCallback = Callable[[PortForwardEvent], None]
PortForwardChannelFactory = Callable[[], Awaitable[int]]


def _detect_protocol(sample: bytes) -> Optional[str]:
Expand Down Expand Up @@ -226,13 +227,14 @@ async def run_remote_file(

async def run_port_forward(
m2m: M2MClient,
channel_port: int,
channel_port: Optional[int],
local_port: int,
*,
channel_factory: Optional[PortForwardChannelFactory] = None,
event_callback: Optional[PortForwardEventCallback] = None,
) -> None:
queue = m2m.channel_queue(channel_port)
active_client = asyncio.Lock()
initial_channel_claimed = False
initial_channel_lock = asyncio.Lock()

def emit(kind: str, *, bytes_count: int = 0, detail: str = "") -> None:
if event_callback:
Expand All @@ -245,29 +247,35 @@ def emit(kind: str, *, bytes_count: int = 0, detail: str = "") -> None:
)
)

async def allocate_channel() -> int:
nonlocal initial_channel_claimed
async with initial_channel_lock:
if channel_port is not None and not initial_channel_claimed:
initial_channel_claimed = True
return channel_port
if channel_factory is None:
raise RuntimeError("No remote channel available for this connection.")
return await channel_factory()

async def handle_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
peer = writer.get_extra_info("peername")
peer_label = str(peer) if peer else "unknown"
if active_client.locked():
emit("connection_rejected", detail=peer_label)
writer.close()
await writer.wait_closed()
return

await active_client.acquire()
acquired = True
emit("connection_opened", detail=peer_label)
connection_label = f"{peer_label}/{id(writer):x}"
channel_for_client: Optional[int] = None
emit("connection_opened", detail=connection_label)
local_started = asyncio.Event()
first_remote_byte_seen = False
protocols_seen = set()

async def local_to_remote() -> None:
if channel_for_client is None:
return
while True:
data = await reader.read(65536)
if not data:
break
local_started.set()
await m2m.send_route(channel_port, data)
await m2m.send_route(channel_for_client, data)
emit("bytes_up", bytes_count=len(data))
protocol = _detect_protocol(data[:32])
if protocol and protocol not in protocols_seen:
Expand All @@ -276,13 +284,16 @@ async def local_to_remote() -> None:

async def remote_to_local() -> None:
nonlocal first_remote_byte_seen
if channel_for_client is None:
return
queue = m2m.channel_queue(channel_for_client)
while True:
data = await queue.get()
if data is None:
break
if local_started.is_set() and not first_remote_byte_seen:
first_remote_byte_seen = True
emit("first_remote_byte")
emit("first_remote_byte", detail=connection_label)
writer.write(data)
await writer.drain()
emit("bytes_down", bytes_count=len(data))
Expand All @@ -292,13 +303,29 @@ async def remote_to_local() -> None:
emit("protocol_detected", detail=protocol)

try:
await asyncio.gather(local_to_remote(), remote_to_local())
try:
channel_for_client = await allocate_channel()
except Exception as exc:
emit("connection_rejected", detail=f"{connection_label}: {exc}")
return
to_remote_task = asyncio.create_task(local_to_remote())
to_local_task = asyncio.create_task(remote_to_local())
done, pending = await asyncio.wait(
{to_remote_task, to_local_task},
return_when=asyncio.FIRST_COMPLETED,
)
for task in pending:
task.cancel()
await asyncio.gather(*done, *pending, return_exceptions=True)
finally:
if channel_for_client is not None:
try:
await m2m.close_channel(channel_for_client)
except Exception as exc:
emit("channel_close_failed", detail=f"{connection_label}: {exc}")
writer.close()
await writer.wait_closed()
emit("connection_closed", detail=peer_label)
if acquired and active_client.locked():
active_client.release()
emit("connection_closed", detail=connection_label)

server = await asyncio.start_server(handle_client, host="127.0.0.1", port=local_port)
emit("listener_started", detail=f"127.0.0.1:{local_port}")
Expand Down
29 changes: 29 additions & 0 deletions tests/test_devices_history_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -223,6 +223,35 @@ def fake_request(
self.assertEqual(payload.get("id"), 777)
self.assertEqual(payload.get("is_deleted"), True)

def test_history_list_debug_prints_endpoint(self) -> None:
timeline_payload = {"results": [], "count": 0}

def fake_get(_self, path: str, *, params=None): # type: ignore[no-untyped-def]
_ = params
self.assertEqual(path, f"/api/developer/devices/{self.device_hash}/timeline/")
return ApiResponse(True, 200, timeline_payload, json.dumps(timeline_payload))

with tempfile.TemporaryDirectory() as temp_dir:
config_path = Path(temp_dir) / "cli.json"
self._write_authed_config(config_path)
with patch("dataplicity_cli.cli.ApiClient.get", new=fake_get):
result = self.runner.invoke(
app,
[
"--debug",
"--config",
str(config_path),
"devices",
"history",
"list",
self.device_hash,
],
)

self.assertEqual(result.exit_code, 0, msg=result.output)
self.assertIn("/api/developer/devices/", result.output)
self.assertIn("/timeline/", result.output)


if __name__ == "__main__":
unittest.main()
Loading
Loading