1"""Remote side of a portal session.
2
3Runs on the production machine. Connects to the relay, prints a portal
4code, waits for the local side to connect, then executes commands as
5they arrive through the encrypted tunnel.
6"""
7
8from __future__ import annotations
9
10import ast
11import asyncio
12import base64
13import contextlib
14import json
15import os
16import sys
17import traceback
18from contextlib import redirect_stderr, redirect_stdout
19from datetime import datetime
20
21from websockets.asyncio.client import ClientConnection
22from websockets.asyncio.client import connect as ws_connect
23from websockets.exceptions import ConnectionClosed
24
25from .codegen import generate_code
26from .crypto import PortalEncryptor, channel_id, perform_key_exchange
27from .protocol import (
28 DEFAULT_EXEC_TIMEOUT,
29 DEFAULT_RELAY_HOST,
30 FILE_CHUNK_SIZE,
31 MAX_FILE_SIZE,
32 chunk_count,
33 make_error,
34 make_exec_result,
35 make_exec_stdout,
36 make_file_data,
37 make_file_push_result,
38 make_ping,
39 make_pong,
40 make_relay_url,
41)
42
43_real_stdout = sys.stdout
44
45
46def _log(msg: str) -> None:
47 ts = datetime.now().strftime("%H:%M:%S")
48 _real_stdout.write(f"[{ts}] {msg}\n")
49 _real_stdout.flush()
50
51
52async def _send_error(
53 ws: ClientConnection,
54 encryptor: PortalEncryptor,
55 req_id: int | None,
56 error_text: str,
57) -> None:
58 """Send an error response back through the tunnel."""
59 msg = make_error(error_text)
60 msg["_req_id"] = req_id
61 await ws.send(encryptor.encrypt_message(msg))
62
63
64class _TunnelWriter:
65 """File-like that streams writes through the tunnel as exec_stdout messages."""
66
67 def __init__(
68 self,
69 loop: asyncio.AbstractEventLoop,
70 ws: ClientConnection,
71 encryptor: PortalEncryptor,
72 req_id: int | None,
73 ) -> None:
74 self._loop = loop
75 self._ws = ws
76 self._encryptor = encryptor
77 self._req_id = req_id
78 self._buffer = ""
79
80 def write(self, s: str) -> int:
81 self._buffer += s
82 while "\n" in self._buffer:
83 line, self._buffer = self._buffer.split("\n", 1)
84 self._send(line + "\n")
85 return len(s)
86
87 def flush(self) -> None:
88 if self._buffer:
89 self._send(self._buffer)
90 self._buffer = ""
91
92 def _send(self, data: str) -> None:
93 msg = make_exec_stdout(data)
94 msg["_req_id"] = self._req_id
95 future = asyncio.run_coroutine_threadsafe(
96 self._ws.send(self._encryptor.encrypt_message(msg)),
97 self._loop,
98 )
99 try:
100 future.result(timeout=30)
101 except Exception:
102 pass # Don't crash exec for a send failure
103
104
105async def run_remote(
106 *,
107 writable: bool = False,
108 timeout_minutes: int = 30,
109 relay_host: str = DEFAULT_RELAY_HOST,
110) -> None:
111 """Start the remote side of a portal session."""
112
113 code = generate_code()
114
115 mode = "writable" if writable else "read-only"
116 print(f"Portal code: {code}")
117 print(f"Session mode: {mode}")
118 print("Waiting for connection...")
119 print()
120
121 cid = channel_id(code)
122 relay_url = make_relay_url(relay_host, cid, "start")
123
124 max_output = (
125 1024 * 1024
126 ) # 1MB — truncate return values to prevent massive relay payloads
127 tmp_prefix = os.path.realpath("/tmp") # Resolve once (macOS: /tmp → /private/tmp)
128
129 def execute_code(
130 code_str: str,
131 *,
132 json_output: bool = False,
133 output_writer: _TunnelWriter,
134 ) -> dict:
135 """Execute Python code, streaming stdout through the tunnel.
136
137 Each execution gets a fresh namespace. The last expression's value
138 is captured as the return value (like the interactive REPL).
139 """
140 namespace: dict = {}
141 return_value = None
142 error = None
143
144 try:
145 tree = ast.parse(code_str, mode="exec")
146
147 last_expr: ast.Expr | None = None
148 if tree.body and isinstance(tree.body[-1], ast.Expr):
149 popped = tree.body.pop()
150 assert isinstance(popped, ast.Expr)
151 last_expr = popped
152
153 # Process-global redirect — safe because _log() uses _real_stdout
154 ctx = contextlib.ExitStack()
155 ctx.enter_context(redirect_stdout(output_writer))
156 ctx.enter_context(redirect_stderr(output_writer))
157 if not writable:
158 try:
159 from plain.postgres.db import read_only
160
161 ctx.enter_context(read_only())
162 except Exception:
163 pass # No DB configured or plain-postgres not installed
164
165 with ctx:
166 if tree.body:
167 compiled = compile(tree, "<portal>", "exec")
168 exec(compiled, namespace) # noqa: S102
169
170 if last_expr is not None:
171 expr_code = compile(
172 ast.Expression(last_expr.value), "<portal>", "eval"
173 )
174 result = eval(expr_code, namespace) # noqa: S307
175 if result is not None:
176 if json_output:
177 try:
178 return_value = json.dumps(result)
179 except (TypeError, ValueError):
180 return_value = repr(result)
181 else:
182 return_value = repr(result)
183
184 except BaseException:
185 error = traceback.format_exc()
186 finally:
187 # Flush any remaining buffered output
188 output_writer.flush()
189 # Close DB connection to prevent leaks across to_thread calls
190 try:
191 from plain.postgres.db import get_connection, has_connection
192
193 if has_connection():
194 get_connection().close()
195 except Exception:
196 pass
197
198 if return_value and len(return_value) > max_output:
199 return_value = (
200 return_value[:max_output]
201 + f"\n... truncated ({len(return_value)} bytes total)"
202 )
203
204 return {
205 "return_value": return_value,
206 "error": error,
207 }
208
209 async def handle_file_pull(remote_path: str, req_id: int | None) -> None:
210 """Read a file from disk and send it in chunks."""
211 try:
212 file_size = os.path.getsize(remote_path)
213 if file_size > MAX_FILE_SIZE:
214 await _send_error(
215 ws,
216 encryptor,
217 req_id,
218 f"File too large: {file_size} bytes (max {MAX_FILE_SIZE})",
219 )
220 return
221
222 name = os.path.basename(remote_path)
223 chunks = chunk_count(file_size)
224
225 _log(f" sending {name} ({file_size} bytes, {chunks} chunks)")
226
227 with open(remote_path, "rb") as f:
228 for i in range(chunks):
229 data = f.read(FILE_CHUNK_SIZE)
230 msg = make_file_data(name=name, chunk=i, chunks=chunks, data=data)
231 msg["_req_id"] = req_id
232 await ws.send(encryptor.encrypt_message(msg))
233
234 except FileNotFoundError:
235 await _send_error(ws, encryptor, req_id, f"File not found: {remote_path}")
236 except (PermissionError, IsADirectoryError, OSError) as e:
237 await _send_error(ws, encryptor, req_id, f"{type(e).__name__}: {e}")
238
239 async def handle_file_push(msg: dict) -> None:
240 """Receive a file chunk and write it to disk."""
241 req_id = msg.get("_req_id")
242 remote_path = msg["remote_path"]
243 chunk_idx = msg["chunk"]
244 chunks = msg["chunks"]
245 data = base64.b64decode(msg["data"])
246
247 resolved = os.path.realpath(remote_path)
248 if not resolved.startswith(tmp_prefix + "/"):
249 await _send_error(
250 ws,
251 encryptor,
252 req_id,
253 f"Push restricted to /tmp/. Got: {remote_path} (resolved: {resolved})",
254 )
255 return
256
257 if chunk_idx == 0:
258 _log(f"push: {remote_path} ({chunks} chunks)")
259
260 try:
261 mode = "wb" if chunk_idx == 0 else "ab"
262 with open(remote_path, mode) as f:
263 f.write(data)
264 except OSError as e:
265 await _send_error(ws, encryptor, req_id, f"{type(e).__name__}: {e}")
266 return
267
268 # Ack every chunk so the sender doesn't block waiting
269 if chunk_idx == chunks - 1:
270 total_bytes = os.path.getsize(remote_path)
271 _log(f" received {total_bytes} bytes")
272 result = make_file_push_result(path=remote_path, total_bytes=total_bytes)
273 else:
274 result = {"type": "file_push_ack", "chunk": chunk_idx}
275 result["_req_id"] = req_id
276 await ws.send(encryptor.encrypt_message(result))
277
278 async with ws_connect(relay_url) as ws:
279 encryptor = await perform_key_exchange(ws, code, side="start")
280 _log("Connected from remote client.")
281
282 last_activity = asyncio.get_running_loop().time()
283
284 async def check_timeout() -> None:
285 nonlocal last_activity
286 if timeout_minutes <= 0:
287 return
288 while True:
289 await asyncio.sleep(60)
290 idle = asyncio.get_running_loop().time() - last_activity
291 remaining = (timeout_minutes * 60) - idle
292 if remaining <= 60 and remaining > 0:
293 print(
294 f"\nWarning: session will timeout in {int(remaining)} seconds due to inactivity.",
295 flush=True,
296 )
297 if idle >= timeout_minutes * 60:
298 print(
299 "\nSession timed out due to inactivity.",
300 flush=True,
301 )
302 await ws.close()
303 return
304
305 timeout_task = asyncio.create_task(check_timeout())
306
307 async def send_keepalive_pings() -> None:
308 while True:
309 await asyncio.sleep(30)
310 await ws.send(encryptor.encrypt_message(make_ping()))
311
312 keepalive_task = asyncio.create_task(send_keepalive_pings())
313
314 try:
315 async for raw in ws:
316 last_activity = asyncio.get_running_loop().time()
317
318 if isinstance(raw, str):
319 continue
320
321 msg = encryptor.decrypt_message(raw)
322 msg_type = msg.get("type")
323
324 if msg_type == "ping":
325 await ws.send(encryptor.encrypt_message(make_pong()))
326
327 elif msg_type == "pong":
328 pass
329
330 elif msg_type == "exec":
331 req_id = msg.get("_req_id")
332 code_str = msg["code"]
333 json_output = msg.get("json_output", False)
334 exec_timeout = msg.get("timeout", DEFAULT_EXEC_TIMEOUT)
335 _log(
336 f"exec: {code_str[:200]}{'...' if len(code_str) > 200 else ''}"
337 )
338 # Create a writer that streams stdout through the tunnel
339 tunnel_writer = _TunnelWriter(
340 asyncio.get_running_loop(), ws, encryptor, req_id
341 )
342 try:
343 result = await asyncio.wait_for(
344 asyncio.to_thread(
345 execute_code,
346 code_str,
347 json_output=json_output,
348 output_writer=tunnel_writer,
349 ),
350 timeout=exec_timeout,
351 )
352 except TimeoutError:
353 result = {
354 "return_value": None,
355 "error": f"Execution timed out ({exec_timeout} seconds). The code may still be running in the background.",
356 }
357 return_value = result.get("return_value")
358 error = result.get("error")
359 display = return_value or error or ""
360 if display:
361 _log(
362 f" → {display[:200]}{'...' if len(display) > 200 else ''}"
363 )
364 # Send final result — stdout was already streamed
365 response = make_exec_result(
366 return_value=return_value,
367 error=error,
368 )
369 response["_req_id"] = req_id
370 await ws.send(encryptor.encrypt_message(response))
371
372 elif msg_type == "file_pull":
373 req_id = msg.get("_req_id")
374 remote_path = msg["remote_path"]
375 _log(f"pull: {remote_path}")
376 await handle_file_pull(remote_path, req_id)
377
378 elif msg_type == "file_push":
379 await handle_file_push(msg)
380
381 else:
382 _log(f"Unknown message type: {msg_type}")
383
384 except ConnectionClosed:
385 pass # Normal when relay or network drops the connection
386 finally:
387 timeout_task.cancel()
388 keepalive_task.cancel()
389
390 _log("Client disconnected.")