prime sandbox ssh

examples/sandbox_port_expose_demo.py

@@ -0,0 +1,195 @@
+import socket
+import time
+import urllib.request
+
+from prime_sandboxes import APIClient, APIError, CreateSandboxRequest, SandboxClient
+
+
+def verify_http(url: str) -> bool:
+    """Verify HTTP endpoint is accessible and returns expected response."""
+    try:
+        # Add User-Agent header to avoid 403 from bot protection
+        req = urllib.request.Request(url, headers={"User-Agent": "curl/8.0"})
+        with urllib.request.urlopen(req, timeout=10) as response:
+            status = response.getcode()
+            body = response.read().decode("utf-8")
+            # Python's http.server returns a directory listing HTML
+            if status == 200 and "Directory listing" in body:
+                return True
+            return False
+    except Exception as e:
+        print(f"    HTTP verification error: {e}")
+        return False
+
+
+def verify_tcp(endpoint: str, test_message: bytes = b"Hello") -> bool:
+    """Verify TCP endpoint is accessible and echoes back data."""
+    try:
+        # Parse host:port from endpoint address
+        host, port_str = endpoint.rsplit(":", 1)
+        port = int(port_str)
+
+        # Connect with raw TCP
+        with socket.create_connection((host, port), timeout=10) as sock:
+            sock.sendall(test_message)
+            response = sock.recv(1024)
+            expected = b"Echo: " + test_message
+            return response == expected
+    except Exception as e:
+        print(f"    TCP verification error: {e}")
+        return False
+
+
+def main() -> None:
+    """Demonstrate HTTP and TCP port exposure"""
+    try:
+        client = APIClient()
+        sandbox_client = SandboxClient(client)
+
+        request = CreateSandboxRequest(
+            name="port-expose-demo",
+            docker_image="python:3.11-slim",
+            start_command="tail -f /dev/null",
+            cpu_cores=1,
+            memory_gb=2,
+            timeout_minutes=30,
+        )
+
+        print("Creating sandbox...")
+        sandbox = sandbox_client.create(request)
+        print(f"Created: {sandbox.name} ({sandbox.id})")
+
+        print("\nWaiting for sandbox to be running...")
+        sandbox_client.wait_for_creation(sandbox.id, max_attempts=60)
+        print("Sandbox is running!")
+
+        print("\n--- HTTP Port Exposure ---")
+        print("Starting HTTP server on port 8000...")
+        sandbox_client.execute_command(
+            sandbox.id,
+            "nohup python -m http.server 8000 > /tmp/http.log 2>&1 &",
+        )
+        time.sleep(2)  # Give server time to start
+
+        # Expose the HTTP port
+        http_exposure = sandbox_client.expose(
+            sandbox_id=sandbox.id,
+            port=8000,
+            name="web-server",
+            protocol="HTTP",
+        )
+        print("HTTP port exposed!")
+        print(f"  Exposure ID: {http_exposure.exposure_id}")
+        print(f"  URL: {http_exposure.url}")
+        print(f"  TLS Socket: {http_exposure.tls_socket}")
+        time.sleep(10)
+
+        # Verify HTTP endpoint is accessible
+        print("  Verifying HTTP endpoint...")
+        if verify_http(http_exposure.url):
+            print("  HTTP verification: SUCCESS")
+        else:
+            print("  HTTP verification: FAILED")
+
+        # Start a TCP echo server in the sandbox
+        print("\n--- TCP Port Exposure ---")
+        print("Starting TCP echo server on port 9000...")
+
+        # Create a simple TCP echo server
+        tcp_server_code = """
+import socket
+import threading
+
+def handle_client(conn, addr):
+    print(f"Connection from {addr}")
+    while True:
+        data = conn.recv(1024)
+        if not data:
+            break
+        conn.sendall(b"Echo: " + data)
+    conn.close()
+
+server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+server.bind(("0.0.0.0", 9000))
+server.listen(5)
+print("TCP server listening on port 9000")
+
+while True:
+    conn, addr = server.accept()
+    thread = threading.Thread(target=handle_client, args=(conn, addr))
+    thread.daemon = True
+    thread.start()
+"""
+        # Write and run the TCP server
+        sandbox_client.execute_command(
+            sandbox.id,
+            f"cat > /tmp/tcp_server.py << 'SCRIPT'\n{tcp_server_code}\nSCRIPT",
+        )
+        sandbox_client.execute_command(
+            sandbox.id,
+            "nohup python /tmp/tcp_server.py > /tmp/tcp.log 2>&1 &",
+        )
+        time.sleep(2)  # Give server time to start
+
+        # Expose the TCP port
+        tcp_exposure = sandbox_client.expose(
+            sandbox_id=sandbox.id,
+            port=9000,
+            name="echo-server",
+            protocol="TCP",
+        )
+        print("TCP port exposed!")
+        print(f"  Exposure ID: {tcp_exposure.exposure_id}")
+        print(f"  External Endpoint: {tcp_exposure.external_endpoint}")
+        if tcp_exposure.external_port:
+            print(f"  External Port: {tcp_exposure.external_port}")
+        time.sleep(120)
+
+        # Verify TCP endpoint is accessible
+        print("  Verifying TCP endpoint...")
+        if verify_tcp(tcp_exposure.external_endpoint):
+            print("  TCP verification: SUCCESS (echo server responded correctly)")
+        else:
+            print("  TCP verification: FAILED")
+
+        # List all exposed ports
+        print("\n--- All Exposed Ports ---")
+        ports_response = sandbox_client.list_exposed_ports(sandbox.id)
+        for port in ports_response.exposures:
+            print(f"  {port.name} (port {port.port}):")
+            print(f"    Protocol: {port.protocol}")
+            print(f"    Exposure ID: {port.exposure_id}")
+            if port.protocol == "HTTP":
+                print(f"    URL: {port.url}")
+            else:
+                print(f"    External Endpoint: {port.external_endpoint}")
+
+        # Usage instructions
+        print("\n--- How to Connect ---")
+        print(f"HTTP: curl {http_exposure.url}")
+        print(f"TCP:  Connect to {tcp_exposure.external_endpoint} with a TCP client")
+
+        # Clean up exposures
+        print("\n--- Cleanup ---")
+        print("Removing port exposures...")
+        sandbox_client.unexpose(sandbox.id, http_exposure.exposure_id)
+        print(f"  Removed HTTP exposure: {http_exposure.exposure_id}")
+        sandbox_client.unexpose(sandbox.id, tcp_exposure.exposure_id)
+        print(f"  Removed TCP exposure: {tcp_exposure.exposure_id}")
+
+        # Delete sandbox
+        print(f"\nDeleting sandbox {sandbox.name}...")
+        sandbox_client.delete(sandbox.id)
+        print("Done!")
+
+    except APIError as e:
+        print(f"API Error: {e}")
+        print("Make sure you're logged in: run 'prime login' first")
+    except Exception as e:
+        print(f"Error: {e}")
+        raise
+
+
+if __name__ == "__main__":
+    main()

packages/prime-sandboxes/src/prime_sandboxes/__init__.py

@@ -37,6 +37,7 @@
     Sandbox,
     SandboxListResponse,
     SandboxStatus,
+    SSHSession,
     UpdateSandboxRequest,
 )
 from .sandbox import AsyncSandboxClient, AsyncTemplateClient, SandboxClient, TemplateClient
@@ -76,6 +77,7 @@
     "ExposePortRequest",
     "ExposedPort",
     "ListExposedPortsResponse",
+    "SSHSession",
     # Exceptions
     "APIError",
     "UnauthorizedError",

packages/prime-sandboxes/src/prime_sandboxes/models.py

@@ -178,6 +178,7 @@ class ExposePortRequest(BaseModel):
 
     port: int
     name: Optional[str] = None
+    protocol: str = "HTTP"  # HTTP or TCP/UDP
 
 
 class ExposedPort(BaseModel):
@@ -190,6 +191,8 @@ class ExposedPort(BaseModel):
     url: str
     tls_socket: str
     protocol: Optional[str] = None
+    external_port: Optional[int] = None  # For TCP/UDP exposures
+    external_endpoint: Optional[str] = None  # For TCP/UDP: host:port endpoint
     created_at: Optional[str] = None
 
 
@@ -199,6 +202,23 @@ class ListExposedPortsResponse(BaseModel):
     exposures: List[ExposedPort]
 
 
+class SSHSession(BaseModel):
+    """SSH session details"""
+
+    session_id: str
+    exposure_id: str
+    sandbox_id: str
+    host: str
+    port: int
+    external_endpoint: str
+    expires_at: datetime
+    ttl_seconds: int
+    gateway_url: str
+    user_ns: str
+    job_id: str
+    token: str
+
+
 class BackgroundJob(BaseModel):
     """Background job handle returned when starting a background job"""
 

packages/prime-sandboxes/src/prime_sandboxes/sandbox.py

@@ -43,6 +43,7 @@
     Sandbox,
     SandboxListResponse,
     SandboxLogsResponse,
+    SSHSession,
 )
 
 # Retry configuration for transient connection errors on gateway requests
@@ -713,9 +714,10 @@ def expose(
         sandbox_id: str,
         port: int,
         name: Optional[str] = None,
+        protocol: str = "HTTP",
     ) -> ExposedPort:
-        """Expose an HTTP port from a sandbox."""
-        request = ExposePortRequest(port=port, name=name)
+        """Expose a port from a sandbox."""
+        request = ExposePortRequest(port=port, name=name, protocol=protocol)
         response = self.client.request(
             "POST",
             f"/sandbox/{sandbox_id}/expose",
@@ -732,6 +734,31 @@ def list_exposed_ports(self, sandbox_id: str) -> ListExposedPortsResponse:
         response = self.client.request("GET", f"/sandbox/{sandbox_id}/expose")
         return ListExposedPortsResponse.model_validate(response)
 
+    def list_all_exposed_ports(self) -> ListExposedPortsResponse:
+        """List all exposed ports across all sandboxes for the current user"""
+        response = self.client.request("GET", "/sandbox/expose/all")
+        return ListExposedPortsResponse.model_validate(response)
+
+    def create_ssh_session(
+        self,
+        sandbox_id: str,
+        ttl_seconds: Optional[int] = None,
+    ) -> SSHSession:
+        """Create an SSH session"""
+        payload: Dict[str, Any] = {}
+        if ttl_seconds is not None:
+            payload["ttl_seconds"] = ttl_seconds
+        response = self.client.request(
+            "POST",
+            f"/sandbox/{sandbox_id}/ssh-session",
+            json=payload,
+        )
+        return SSHSession.model_validate(response)
+
+    def close_ssh_session(self, sandbox_id: str, session_id: str) -> None:
+        """Close an SSH session and remove its exposure"""
+        self.client.request("DELETE", f"/sandbox/{sandbox_id}/ssh-session/{session_id}")
+
 
 class AsyncSandboxClient:
     """Async client for sandbox API operations"""
@@ -1293,9 +1320,10 @@ async def expose(
         sandbox_id: str,
         port: int,
         name: Optional[str] = None,
+        protocol: str = "HTTP",
     ) -> ExposedPort:
-        """Expose an HTTP port from a sandbox."""
-        request = ExposePortRequest(port=port, name=name)
+        """Expose a port from a sandbox."""
+        request = ExposePortRequest(port=port, name=name, protocol=protocol)
         response = await self.client.request(
             "POST",
             f"/sandbox/{sandbox_id}/expose",
@@ -1312,6 +1340,31 @@ async def list_exposed_ports(self, sandbox_id: str) -> ListExposedPortsResponse:
         response = await self.client.request("GET", f"/sandbox/{sandbox_id}/expose")
         return ListExposedPortsResponse.model_validate(response)
 
+    async def list_all_exposed_ports(self) -> ListExposedPortsResponse:
+        """List all exposed ports across all sandboxes for the current user"""
+        response = await self.client.request("GET", "/sandbox/expose/all")
+        return ListExposedPortsResponse.model_validate(response)
+
+    async def create_ssh_session(
+        self,
+        sandbox_id: str,
+        ttl_seconds: Optional[int] = None,
+    ) -> SSHSession:
+        """Create an SSH session"""
+        payload: Dict[str, Any] = {}
+        if ttl_seconds is not None:
+            payload["ttl_seconds"] = ttl_seconds
+        response = await self.client.request(
+            "POST",
+            f"/sandbox/{sandbox_id}/ssh-session",
+            json=payload,
+        )
+        return SSHSession.model_validate(response)
+
+    async def close_ssh_session(self, sandbox_id: str, session_id: str) -> None:
+        """Close an SSH session and remove its exposure"""
+        await self.client.request("DELETE", f"/sandbox/{sandbox_id}/ssh-session/{session_id}")
+
 
 class TemplateClient:
     """Client for template/registry helper APIs."""