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