1from __future__ import annotations
  2
  3import asyncio
  4import base64
  5import json
  6import os
  7import sys
  8
  9import click
 10
 11from plain.cli import register_cli
 12
 13from .protocol import (
 14    DEFAULT_EXEC_TIMEOUT,
 15    DEFAULT_RELAY_HOST,
 16    FILE_CHUNK_SIZE,
 17    MAX_FILE_SIZE,
 18    chunk_count,
 19    make_exec,
 20    make_file_pull,
 21    make_file_push,
 22)
 23
 24
 25def _check_response(response: dict) -> None:
 26    """Exit with an error message if the response indicates failure."""
 27    error = response.get("error")
 28    if error:
 29        print(error, file=sys.stderr)
 30        sys.exit(1)
 31
 32
 33@register_cli("portal")
 34@click.group()
 35def cli() -> None:
 36    """Remote Python shell and file transfer via encrypted tunnel."""
 37
 38
 39@cli.command()
 40@click.option(
 41    "--writable", is_flag=True, help="Allow database writes (default: read-only)."
 42)
 43@click.option(
 44    "--timeout",
 45    default=30,
 46    type=int,
 47    help="Idle timeout in minutes (0 to disable).",
 48)
 49@click.option(
 50    "--relay-host",
 51    envvar="PLAIN_PORTAL_RELAY_HOST",
 52    default=DEFAULT_RELAY_HOST,
 53    hidden=True,
 54)
 55@click.option("--yes", "-y", is_flag=True, help="Skip confirmation prompt.")
 56def start(writable: bool, timeout: int, relay_host: str, yes: bool) -> None:
 57    """Start a portal session on the remote machine."""
 58    if writable and not yes:
 59        if not click.confirm(
 60            "This session allows writes to the production database. Continue?"
 61        ):
 62            return
 63
 64    from .remote import run_remote
 65
 66    asyncio.run(
 67        run_remote(writable=writable, timeout_minutes=timeout, relay_host=relay_host)
 68    )
 69
 70
 71@cli.command()
 72@click.argument("code")
 73@click.option(
 74    "--relay-host",
 75    envvar="PLAIN_PORTAL_RELAY_HOST",
 76    default=DEFAULT_RELAY_HOST,
 77    hidden=True,
 78)
 79def connect(code: str, relay_host: str) -> None:
 80    """Connect to a remote portal session."""
 81    from .local import connect as do_connect
 82
 83    asyncio.run(do_connect(code, relay_host=relay_host))
 84
 85
 86@cli.command("exec")
 87@click.argument("code")
 88@click.option("--json", "json_output", is_flag=True, help="Output as JSON.")
 89@click.option(
 90    "--timeout",
 91    default=DEFAULT_EXEC_TIMEOUT,
 92    type=int,
 93    help=f"Execution timeout in seconds (default: {DEFAULT_EXEC_TIMEOUT}).",
 94)
 95def exec_command(code: str, json_output: bool, timeout: int) -> None:
 96    """Execute Python code on the remote machine."""
 97    from .local import send_exec_streaming
 98
 99    request = make_exec(code, json_output=json_output, timeout=timeout)
100
101    # Collect streamed stdout for --json mode, print directly otherwise
102    stdout_parts: list[str] = []
103
104    def on_stdout(data: str) -> None:
105        if json_output:
106            stdout_parts.append(data)
107        else:
108            print(data, end="", flush=True)
109
110    response = asyncio.run(send_exec_streaming(request, on_stdout))
111    _check_response(response)
112
113    if json_output:
114        print(
115            json.dumps(
116                {
117                    "stdout": "".join(stdout_parts),
118                    "return_value": response.get("return_value"),
119                    "error": response.get("error"),
120                }
121            )
122        )
123    else:
124        return_value = response.get("return_value")
125        if return_value is not None:
126            print(f"→ {return_value}")
127
128
129@cli.command()
130@click.argument("remote_path")
131@click.argument("local_path")
132def pull(remote_path: str, local_path: str) -> None:
133    """Pull a file from the remote machine."""
134    from .local import send_command
135
136    request = make_file_pull(remote_path)
137    response = asyncio.run(send_command(request))
138    _check_response(response)
139
140    if response.get("type") == "file_data":
141        data = base64.b64decode(response["data"])
142        with open(local_path, "wb") as f:
143            f.write(data)
144        print(f"Pulled {remote_path} → {local_path} ({len(data)} bytes)")
145    else:
146        print(f"Unexpected response: {response}", file=sys.stderr)
147        sys.exit(1)
148
149
150@cli.command()
151@click.argument("local_path")
152@click.argument("remote_path")
153def push(local_path: str, remote_path: str) -> None:
154    """Push a file to the remote machine."""
155    from .local import send_command
156
157    if not os.path.exists(local_path):
158        print(f"File not found: {local_path}", file=sys.stderr)
159        sys.exit(1)
160
161    file_size = os.path.getsize(local_path)
162    if file_size > MAX_FILE_SIZE:
163        print(
164            f"File too large: {file_size} bytes (max {MAX_FILE_SIZE})",
165            file=sys.stderr,
166        )
167        sys.exit(1)
168
169    async def _push_all() -> dict:
170        chunks = chunk_count(file_size)
171        response = {}
172        with open(local_path, "rb") as f:
173            for i in range(chunks):
174                data = f.read(FILE_CHUNK_SIZE)
175                request = make_file_push(
176                    remote_path=remote_path, chunk=i, chunks=chunks, data=data
177                )
178                response = await send_command(request)
179                if response.get("error"):
180                    return response
181        return response
182
183    response = asyncio.run(_push_all())
184    _check_response(response)
185
186    total_bytes = response.get("bytes", file_size)
187    print(f"Pushed {local_path} → {remote_path} ({total_bytes} bytes)")