boxd 0.1.0.dev2__py3-none-any.whl

boxd/auth.py

@@ -0,0 +1,118 @@
+"""API key authentication — exchange and JWT refresh."""
+
+from __future__ import annotations
+
+import time
+
+import grpc
+import httpx
+
+from .errors import AuthenticationError
+
+
+class TokenAuth:
+    """Manages API key → JWT exchange and transparent refresh."""
+
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+        token: str | None = None,
+        exchange_url: str = "https://boxd.sh/api/v1/auth/token",
+    ):
+        self._api_key = api_key
+        self._token = token
+        self._expires_at: float = 0
+        self._exchange_url = exchange_url
+
+        if not api_key and not token:
+            raise AuthenticationError("provide api_key or token")
+
+    async def get_token(self) -> str:
+        """Return a valid JWT, exchanging/refreshing if needed."""
+        if self._token and not self._api_key:
+            # Direct token mode — no refresh
+            return self._token
+
+        # API key mode — exchange or refresh
+        if not self._token or time.time() > self._expires_at - 300:
+            await self._exchange()
+
+        assert self._token is not None  # set by _exchange()
+        return self._token
+
+    async def _exchange(self) -> None:
+        """Exchange API key for a short-lived JWT."""
+        async with httpx.AsyncClient() as client:
+            resp = await client.post(
+                self._exchange_url,
+                json={"api_key": self._api_key},
+                timeout=10,
+            )
+        if resp.status_code == 401:
+            raise AuthenticationError("invalid or expired API key")
+        if resp.status_code == 429:
+            raise AuthenticationError("rate limit exceeded — try again later")
+        if resp.status_code != 200:
+            raise AuthenticationError(f"exchange failed: {resp.status_code} {resp.text}")
+
+        data = resp.json()
+        self._token = data["token"]
+        self._expires_at = data["expires_at"]
+
+    async def metadata(self) -> list[tuple[str, str]]:
+        """Return auth metadata for manual injection (streaming calls)."""
+        token = await self.get_token()
+        return [("authorization", f"Bearer {token}")]
+
+    def interceptor(self) -> "AuthInterceptor":
+        """Return a gRPC interceptor that injects the Bearer token."""
+        return AuthInterceptor(self)
+
+
+class AuthInterceptor(
+    grpc.aio.UnaryUnaryClientInterceptor,
+    grpc.aio.UnaryStreamClientInterceptor,
+    grpc.aio.StreamUnaryClientInterceptor,
+    grpc.aio.StreamStreamClientInterceptor,
+):
+    """gRPC interceptor that injects Bearer token into metadata."""
+
+    def __init__(self, auth: TokenAuth):
+        self._auth = auth
+
+    async def _add_auth(self, client_call_details):
+        token = await self._auth.get_token()
+        metadata = list(client_call_details.metadata or [])
+        metadata.append(("authorization", f"Bearer {token}"))
+        return client_call_details._replace(metadata=metadata)
+
+    async def _intercept(self, continuation, client_call_details, request_or_iterator):
+        try:
+            new_details = await self._add_auth(client_call_details)
+        except AuthenticationError:
+            raise _abort_unauthenticated()
+        return await continuation(new_details, request_or_iterator)
+
+    async def intercept_unary_unary(self, continuation, client_call_details, request):
+        return await self._intercept(continuation, client_call_details, request)
+
+    async def intercept_unary_stream(self, continuation, client_call_details, request):
+        return await self._intercept(continuation, client_call_details, request)
+
+    async def intercept_stream_unary(self, continuation, client_call_details, request_iterator):
+        return await self._intercept(continuation, client_call_details, request_iterator)
+
+    async def intercept_stream_stream(self, continuation, client_call_details, request_iterator):
+        return await self._intercept(continuation, client_call_details, request_iterator)
+
+
+def _abort_unauthenticated() -> grpc.aio.AioRpcError:
+    """Create a proper gRPC UNAUTHENTICATED error for clean interceptor shutdown."""
+    return grpc.aio.AioRpcError(
+        code=grpc.StatusCode.UNAUTHENTICATED,
+        initial_metadata=grpc.aio.Metadata(),
+        trailing_metadata=grpc.aio.Metadata(),
+        details="authentication failed",
+        debug_error_string=None,
+    )

boxd/box.py

@@ -0,0 +1,320 @@
+"""Box — a running VM with lifecycle, exec, file, and proxy methods."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import grpc.aio
+
+from ._generated import api_pb2
+from ._utils import GrpcCaller
+from .errors import from_grpc_error
+from .exec import ExecProcess, ExecResult, StreamReader, StreamWriter
+from .types import Proxy, ResumeResult, SuspendResult
+
+if TYPE_CHECKING:
+    from .client import Compute
+
+
+@dataclass
+class Box(GrpcCaller):
+    """A boxd VM."""
+
+    id: str
+    name: str
+    image: str
+    public_ip: str
+    status: str
+    _client: Compute = field(repr=False)
+    url: str = ""
+    boot_time_ms: int | None = None
+    #: Source VM id, set when this VM was created via ``compute.box.fork(...)``.
+    forked_from: str | None = None
+    #: Restart policy, populated by ``compute.box.get(...)`` (server doesn't return it on create/fork).
+    restart_policy: str | None = None
+    #: Disk size in bytes, populated by ``compute.box.get(...)``.
+    disk_bytes: int | None = None
+    #: Auto-suspend timeout in seconds (0 = disabled), populated by ``compute.box.get(...)``.
+    auto_suspend_timeout_secs: int | None = None
+
+    # ── Lifecycle ──────────────────────────────────────────────────
+
+    async def destroy(self) -> None:
+        """Destroy this VM."""
+        await self._call("DestroyVm", api_pb2.DestroyVmRequest(vm_id=self.id))
+        self.status = "destroyed"
+
+    async def reboot(self) -> None:
+        """Reboot this VM."""
+        await self._call("RebootVm", api_pb2.RebootVmRequest(vm_id=self.id))
+        self.status = "running"
+
+    async def start(self) -> None:
+        """Start this VM."""
+        await self._call("StartVm", api_pb2.StartVmRequest(vm_id=self.id))
+        self.status = "running"
+
+    async def stop(self) -> None:
+        """Stop this VM."""
+        await self._call("StopVm", api_pb2.StopVmRequest(vm_id=self.id))
+        self.status = "stopped"
+
+    async def suspend(self) -> SuspendResult:
+        """Suspend this VM."""
+        resp = await self._call("SuspendVm", api_pb2.SuspendVmRequest(vm_id=self.id))
+        self.status = "suspended"
+        return SuspendResult(suspend_us=resp.suspend_us)
+
+    async def resume(self) -> ResumeResult:
+        """Resume this VM from suspension."""
+        resp = await self._call("ResumeVm", api_pb2.ResumeVmRequest(vm_id=self.id))
+        self.status = "running"
+        return ResumeResult(resume_us=resp.resume_us)
+
+    # ── Exec ──────────────────────────────────────────────────────
+
+    async def exec(
+        self,
+        *args: str,
+        stream: bool = False,
+        interactive: bool = False,
+        pty: bool = False,
+        text: bool = True,
+        timeout: float | None = None,
+        env: dict[str, str] | None = None,
+    ) -> ExecResult | ExecProcess:
+        """Execute a command in this VM.
+
+        Args:
+            *args: Command and arguments (e.g. ``"ls", "-la"``).
+            stream: If True, return an :class:`ExecProcess` for streaming I/O.
+            interactive: Enable interactive mode (implies pty).
+            pty: Allocate a pseudo-terminal.
+            text: (unused, reserved for future use).
+            timeout: Timeout in seconds for non-streaming mode.
+            env: Extra environment variables for the command.
+        """
+        use_tty = pty or interactive
+        init_chunk = self._build_exec_chunk(list(args), use_tty, env)
+
+        if stream:
+            return await self._exec_stream(init_chunk)
+        return await self._exec_simple(init_chunk, timeout)
+
+    def _build_exec_chunk(
+        self,
+        command: list[str],
+        tty: bool,
+        env: dict[str, str] | None,
+    ) -> api_pb2.ExecChunk:
+        """Build the initial ExecChunk for a command."""
+        import shlex
+
+        cmd_str = " ".join(shlex.quote(arg) for arg in command)
+        if env:
+            env_prefix = " ".join(
+                f"{k}={shlex.quote(v)}" for k, v in env.items()
+            )
+            cmd_str = f"{env_prefix} {cmd_str}"
+        return api_pb2.ExecChunk(vm_id=self.id, command=cmd_str, tty=tty)
+
+    async def _exec_simple(
+        self,
+        init_chunk: api_pb2.ExecChunk,
+        timeout: float | None,
+    ) -> ExecResult:
+        """Run a command and collect all output."""
+        stub = await self._client._ensure_channel()
+
+        async def request_iter():
+            yield init_chunk
+
+        try:
+            metadata = await self._client._auth.metadata()
+            call = stub.Exec(request_iter(), timeout=timeout, metadata=metadata)
+            stdout_parts: list[bytes] = []
+            exit_code = -1
+
+            async for chunk in call:
+                if chunk.data:
+                    stdout_parts.append(chunk.data)
+                exit_code = chunk.exit_code
+
+            raw = b"".join(stdout_parts)
+            return ExecResult(
+                stdout=raw.decode("utf-8", errors="replace"),
+                stderr="",
+                exit_code=exit_code,
+            )
+        except grpc.aio.AioRpcError as e:
+            raise from_grpc_error(e) from e
+
+    async def _exec_stream(
+        self,
+        init_chunk: api_pb2.ExecChunk,
+    ) -> ExecProcess:
+        """Run a command with streaming I/O."""
+        stub = await self._client._ensure_channel()
+
+        stdout_reader = StreamReader()
+        stderr_reader = StreamReader()
+        stdin_writer = StreamWriter()
+
+        process = ExecProcess(
+            stdout=stdout_reader,
+            stderr=stderr_reader,
+            stdin=stdin_writer,
+        )
+
+        async def request_iter():
+            yield init_chunk
+            while True:
+                data = await stdin_writer._get()
+                stdin_writer._task_done()
+                if data is None:
+                    break
+                yield api_pb2.ExecChunk(data=data, stdin=True)
+
+        async def relay(call):
+            exit_code = -1
+            try:
+                async for chunk in call:
+                    if chunk.data:
+                        stdout_reader._feed(chunk.data)
+                    exit_code = chunk.exit_code
+            except grpc.aio.AioRpcError as e:
+                process._set_error(from_grpc_error(e))
+            else:
+                process._set_exit_code(exit_code)
+            finally:
+                stdout_reader._feed_eof()
+                stderr_reader._feed_eof()
+
+        try:
+            metadata = await self._client._auth.metadata()
+            call = stub.Exec(request_iter(), metadata=metadata)
+        except grpc.aio.AioRpcError as e:
+            raise from_grpc_error(e) from e
+
+        process._relay_task = asyncio.create_task(relay(call))
+        return process
+
+    # ── Files ─────────────────────────────────────────────────────
+
+    async def write_file(self, source: str | bytes | Path, dest: str) -> None:
+        """Write a file to the VM.
+
+        Args:
+            source: File content as bytes or str, or a ``Path`` to read from disk.
+            dest: Destination path inside the VM.
+
+        To upload a local file, pass a ``Path`` object::
+
+            await box.write_file(Path("local/file.py"), "/app/file.py")
+        """
+        if isinstance(source, bytes):
+            data = source
+        elif isinstance(source, Path):
+            data = source.read_bytes()
+        else:
+            data = source.encode()
+        await self._call(
+            "UploadFile",
+            api_pb2.UploadFileRequest(vm_id=self.id, path=dest, data=data),
+        )
+
+    async def read_file(self, path: str) -> bytes:
+        """Download a file from the VM.
+
+        Args:
+            path: Path inside the VM.
+
+        Returns:
+            File contents as bytes.
+        """
+        resp = await self._call(
+            "DownloadFile",
+            api_pb2.DownloadFileRequest(vm_id=self.id, path=path),
+        )
+        return resp.data
+
+    # ── Proxies ───────────────────────────────────────────────────
+
+    async def proxies(self) -> list[Proxy]:
+        """List proxies for this VM."""
+        resp = await self._call(
+            "ListProxies",
+            api_pb2.ListProxiesRequest(vm_name=self.name),
+        )
+        return [
+            Proxy(
+                name=p.name,
+                vm_name=p.vm_name,
+                domain=p.domain,
+                port=p.port,
+                is_default=p.is_default,
+            )
+            for p in resp.proxies
+        ]
+
+    async def create_proxy(self, name: str, port: int = 0) -> Proxy:
+        """Create a new proxy subdomain for this VM."""
+        resp = await self._call(
+            "CreateProxy",
+            api_pb2.CreateProxyRequest(name=name, vm_name=self.name, port=port),
+        )
+        return Proxy(
+            name=resp.name,
+            vm_name=resp.vm_name,
+            domain=resp.domain,
+            port=resp.port,
+            is_default=False,
+        )
+
+    async def delete_proxy(self, name: str) -> None:
+        """Delete a proxy subdomain."""
+        await self._call(
+            "DeleteProxy",
+            api_pb2.DeleteProxyRequest(name=name, vm_name=self.name),
+        )
+
+    async def set_proxy_port(self, port: int, name: str = "") -> None:
+        """Change the port for a proxy.
+
+        Args:
+            port: New target port.
+            name: Proxy name (empty string for the default proxy).
+        """
+        await self._call(
+            "SetProxyPort",
+            api_pb2.SetProxyPortRequest(name=name, vm_name=self.name, port=str(port)),
+        )
+
+    # ── Logs ─────────────────────────────────────────────────────
+
+    async def stream_logs(self, follow: bool = False):
+        """Stream console log chunks from the VM.
+
+        Args:
+            follow: If True, keep the stream open and yield new chunks as
+                they arrive. If False, yield only the log content currently
+                available, then end.
+
+        Yields:
+            ``bytes`` chunks of raw log data.
+        """
+        stub = await self._client._ensure_channel()
+        try:
+            metadata = await self._client._auth.metadata()
+            call = stub.StreamLogs(
+                api_pb2.StreamLogsRequest(vm_id=self.id, follow=follow),
+                metadata=metadata,
+            )
+            async for chunk in call:
+                if chunk.data:
+                    yield chunk.data
+        except grpc.aio.AioRpcError as e:
+            raise from_grpc_error(e) from e

boxd/boxes.py

@@ -0,0 +1,164 @@
+"""BoxService — create, list, get, and fork VMs."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import grpc.aio
+
+from ._generated import api_pb2
+from ._utils import GrpcCaller, parse_size
+from .errors import NotFoundError, from_grpc_error
+from .types import BoxConfig
+
+if TYPE_CHECKING:
+    from .box import Box
+    from .client import Compute
+
+
+def _box_config_to_proto(config: BoxConfig | None) -> api_pb2.VmConfig | None:
+    """Convert a BoxConfig dataclass to a VmConfig protobuf message."""
+    if config is None:
+        return None
+
+    vm_config = api_pb2.VmConfig()
+
+    if config.vcpu:
+        vm_config.vcpu = config.vcpu
+    if config.memory:
+        vm_config.memory_bytes = parse_size(config.memory)
+    if config.disk:
+        vm_config.disk_bytes = parse_size(config.disk)
+
+    if config.lifecycle:
+        vm_config.srf.auto_suspend_timeout_secs = config.lifecycle.auto_suspend_timeout
+        vm_config.srf.auto_destroy_timeout_secs = config.lifecycle.auto_destroy_timeout
+
+    if config.network:
+        vm_config.network.ssh = config.network.ssh
+        if config.network.proxies:
+            for p in config.network.proxies:
+                entry = api_pb2.ProxyEntry(name=p.name, port=p.port)
+                vm_config.network.proxies.append(entry)
+
+    if config.volumes:
+        for v in config.volumes:
+            mount = api_pb2.VolumeMount(
+                disk_id=v.disk_id,
+                mount_path=v.mount_path,
+                read_only=v.read_only,
+            )
+            vm_config.volumes.append(mount)
+
+    return vm_config
+
+
+def _vm_response_to_box(resp, client: Compute, *, has_url: bool = False) -> Box:
+    """Convert a VM response proto to a Box dataclass.
+
+    Works with both GetVmResponse (image_ref field) and
+    CreateVmResponse/ForkVmResponse (image field, plus url/boot_time_ms).
+    """
+    from .box import Box
+
+    image = getattr(resp, "image", "") or getattr(resp, "image_ref", "")
+    kwargs: dict = dict(
+        id=resp.vm_id,
+        name=resp.name,
+        image=image,
+        public_ip=resp.public_ip,
+        status=resp.status,
+        _client=client,
+    )
+    if has_url:
+        kwargs["url"] = resp.url
+        kwargs["boot_time_ms"] = resp.boot_time_ms or None
+    # ForkVmResponse carries source vm id; CreateVm/GetVm don't.
+    forked_from = getattr(resp, "forked_from", "")
+    if forked_from:
+        kwargs["forked_from"] = forked_from
+    # GetVmResponse-only fields (it's the only one with image_ref instead of image).
+    if hasattr(resp, "image_ref"):
+        kwargs["restart_policy"] = resp.restart_policy or None
+        kwargs["disk_bytes"] = resp.disk_bytes or None
+        kwargs["auto_suspend_timeout_secs"] = resp.auto_suspend_timeout_secs
+    return Box(**kwargs)
+
+
+class BoxService(GrpcCaller):
+    """Create, list, get, and fork VMs."""
+
+    def __init__(self, client: Compute) -> None:
+        self._client = client
+
+    async def create(
+        self,
+        name: str = "",
+        image: str = "",
+        config: BoxConfig | None = None,
+    ) -> Box:
+        """Create a new VM."""
+        req = api_pb2.CreateVmRequest(name=name)
+        if image:
+            req.image_ref = image
+        if config:
+            if config.env:
+                for k, v in config.env.items():
+                    req.env.append(api_pb2.EnvVar(key=k, value=v))
+            if config.cmd:
+                req.cmd.extend(config.cmd)
+            if config.restart_policy:
+                req.restart_policy = config.restart_policy
+
+            proto_config = _box_config_to_proto(config)
+            if proto_config is not None:
+                req.config.CopyFrom(proto_config)
+
+        resp = await self._call("CreateVm", req)
+        return _vm_response_to_box(resp, self._client, has_url=True)
+
+    async def list(self) -> list[Box]:
+        """List all VMs."""
+        resp = await self._call("ListVms", api_pb2.ListVmsRequest())
+        return [_vm_response_to_box(vm, self._client) for vm in resp.vms]
+
+    async def get(self, vm_id: str) -> Box:
+        """Get a VM by ID or name.
+
+        Tries by ID first; on NotFoundError falls back to name lookup.
+        """
+        stub = await self._client._ensure_channel()
+
+        try:
+            resp = await stub.GetVm(api_pb2.GetVmRequest(vm_id=vm_id))
+            return _vm_response_to_box(resp, self._client)
+        except grpc.aio.AioRpcError as e:
+            err = from_grpc_error(e)
+            if not isinstance(err, NotFoundError):
+                raise err from e
+
+        # Fall back to name lookup via list
+        all_vms = await self.list()
+        for vm in all_vms:
+            if vm.name == vm_id:
+                return vm
+        raise NotFoundError(f"VM not found: {vm_id}")
+
+    async def fork(
+        self,
+        source: str,
+        name: str | None = None,
+        config: BoxConfig | None = None,
+    ) -> Box:
+        """Fork (clone) a VM."""
+        source_vm = await self.get(source)
+        req = api_pb2.ForkVmRequest(source_vm_id=source_vm.id)
+        if name:
+            req.name = name
+
+        proto_config = _box_config_to_proto(config)
+        if proto_config is not None:
+            req.config.CopyFrom(proto_config)
+
+        resp = await self._call("ForkVm", req)
+        return _vm_response_to_box(resp, self._client, has_url=True)