add networking policy interface and image build

Author: Dingway98

hyperbrowser/client/managers/async_manager/sandbox.py

@@ -1,17 +1,28 @@
+import asyncio
+import functools
+import time
 from typing import Dict, Optional, Union
 
 from ....exceptions import HyperbrowserError
 from ....models.sandbox import (
+    CompleteSandboxImageBuildParams,
     CreateSandboxParams,
+    CreateSandboxImageBuildParams,
     SandboxDetail,
     SandboxExecParams,
     SandboxExposeParams,
     SandboxExposeResult,
+    SandboxImageBuild,
+    SandboxImageBuildCreateResult,
+    SandboxImageListParams,
     SandboxImageListResponse,
+    SandboxImageInit,
     SandboxListParams,
     SandboxListResponse,
     SandboxMemorySnapshotParams,
     SandboxMemorySnapshotResult,
+    SandboxNetworkPolicy,
+    SandboxNetworkUpdateResult,
     SandboxRuntimeSession,
     SandboxSnapshotListParams,
     SandboxSnapshotListResponse,
@@ -30,6 +41,15 @@
     _copy_model,
     _expires_within_buffer,
 )
+from ..sandboxes.image_build import (
+    IMAGE_BUILD_SOURCE_PLATFORM,
+    build_docker_image_from_dockerfile,
+    is_terminal_image_build_status,
+    make_temp_docker_tag,
+    package_docker_image,
+    remove_docker_image,
+    upload_image_build_artifact,
+)
 from .sandboxes.sandbox_files import (
     DEFAULT_WATCH_TIMEOUT_MS,
     SandboxFileWatchHandle,
@@ -67,6 +87,12 @@
 ]
 
 
+async def _run_blocking(func, *args, **kwargs):
+    loop = asyncio.get_running_loop()
+    call = functools.partial(func, *args, **kwargs)
+    return await loop.run_in_executor(None, call)
+
+
 class SandboxHandle:
     def __init__(self, service: "SandboxManager", detail: SandboxDetail):
         self._service = service
@@ -130,6 +156,10 @@ def disk_mib(self):
     def exposed_ports(self):
         return self._detail.exposed_ports
 
+    @property
+    def network(self):
+        return self._detail.network
+
     def to_dict(self):
         return self._detail.model_dump()
 
@@ -166,6 +196,25 @@ async def create_memory_snapshot(
             raise TypeError("params must be a SandboxMemorySnapshotParams instance")
         return await self._service.create_memory_snapshot(self.id, normalized)
 
+    async def update_network(
+        self,
+        policy: SandboxNetworkPolicy,
+    ) -> SandboxNetworkUpdateResult:
+        if not isinstance(policy, SandboxNetworkPolicy):
+            raise TypeError("policy must be a SandboxNetworkPolicy instance")
+        result = await self._service.update_network(self.id, policy)
+        self._detail = self._detail.model_copy(update={"network": result.network})
+        return result
+
+    async def clear_network(self) -> SandboxNetworkUpdateResult:
+        return await self.update_network(
+            SandboxNetworkPolicy(
+                allow_internet_access=True,
+                allow_out=[],
+                deny_out=[],
+            )
+        )
+
     async def expose(self, params: SandboxExposeParams) -> SandboxExposeResult:
         if not isinstance(params, SandboxExposeParams):
             raise TypeError("params must be a SandboxExposeParams instance")
@@ -350,8 +399,16 @@ async def list(
         )
         return SandboxListResponse(**payload)
 
-    async def list_images(self) -> SandboxImageListResponse:
-        payload = await self._request("GET", "/images")
+    async def list_images(
+        self, params: SandboxImageListParams = SandboxImageListParams()
+    ) -> SandboxImageListResponse:
+        if not isinstance(params, SandboxImageListParams):
+            raise TypeError("params must be a SandboxImageListParams instance")
+        payload = await self._request(
+            "GET",
+            "/images",
+            params=params.model_dump(exclude_none=True, by_alias=True),
+        )
         return SandboxImageListResponse(**payload)
 
     async def list_snapshots(
@@ -370,6 +427,170 @@ async def stop(self, sandbox_id: str) -> BasicResponse:
         payload = await self._request("PUT", f"/sandbox/{sandbox_id}/stop")
         return BasicResponse(**payload)
 
+    async def create_image_build(
+        self,
+        params: CreateSandboxImageBuildParams,
+    ) -> SandboxImageBuildCreateResult:
+        if not isinstance(params, CreateSandboxImageBuildParams):
+            raise TypeError("params must be a CreateSandboxImageBuildParams instance")
+        payload = await self._request(
+            "POST",
+            "/images/builds",
+            data=params.model_dump(exclude_none=True, by_alias=True),
+        )
+        return SandboxImageBuildCreateResult(**payload)
+
+    async def get_image_build(self, build_id: str) -> SandboxImageBuild:
+        payload = await self._request("GET", f"/images/builds/{build_id}")
+        return SandboxImageBuild(**payload["build"])
+
+    async def complete_image_build(
+        self,
+        build_id: str,
+        params: CompleteSandboxImageBuildParams,
+    ) -> SandboxImageBuild:
+        if not isinstance(params, CompleteSandboxImageBuildParams):
+            raise TypeError("params must be a CompleteSandboxImageBuildParams instance")
+        payload = await self._request(
+            "POST",
+            f"/images/builds/{build_id}/complete",
+            data=params.model_dump(exclude_none=True, by_alias=True),
+        )
+        return SandboxImageBuild(**payload["build"])
+
+    async def cancel_image_build(self, build_id: str) -> SandboxImageBuild:
+        payload = await self._request("POST", f"/images/builds/{build_id}/cancel")
+        return SandboxImageBuild(**payload["build"])
+
+    async def wait_for_image_build(
+        self,
+        build_id: str,
+        *,
+        poll_interval: float = 3.0,
+        timeout: Optional[float] = 35 * 60,
+    ) -> SandboxImageBuild:
+        deadline = None if timeout is None else time.monotonic() + timeout
+        while True:
+            build = await self.get_image_build(build_id)
+            if build.status == "completed":
+                return build
+            if build.status == "failed":
+                message = build.error_message or "image build failed"
+                if build.error_code:
+                    message = f"image build failed [{build.error_code}]: {message}"
+                raise RuntimeError(message)
+            if is_terminal_image_build_status(build.status):
+                raise RuntimeError(f"image build {build.status}")
+            if deadline is not None and time.monotonic() >= deadline:
+                raise TimeoutError(f"timed out waiting for image build {build_id}")
+            await asyncio.sleep(poll_interval)
+
+    async def build_image_from_docker_image(
+        self,
+        *,
+        docker_image: str,
+        image_name: str,
+        platform: str = IMAGE_BUILD_SOURCE_PLATFORM,
+        image_init: Optional[SandboxImageInit] = None,
+        image_config_user: Optional[str] = None,
+        wait: bool = True,
+        poll_interval: float = 3.0,
+        wait_timeout: Optional[float] = 35 * 60,
+        temp_dir: Optional[str] = None,
+        upload_timeout: Optional[float] = None,
+    ) -> SandboxImageBuild:
+        artifact = await _run_blocking(
+            package_docker_image,
+            docker_image,
+            platform=platform,
+            temp_dir=temp_dir,
+        )
+        try:
+            create_result = await self.create_image_build(
+                CreateSandboxImageBuildParams(
+                    image_name=image_name,
+                    input_sha256=artifact.sha256_hex,
+                    input_size_bytes=artifact.size_bytes,
+                    input_format=artifact.input_format,
+                    source_platform=artifact.source_platform,
+                    image_config_user=(
+                        image_config_user
+                        if image_config_user is not None
+                        else artifact.image_config_user
+                    ),
+                    image_init=image_init
+                    if image_init is not None
+                    else artifact.image_init,
+                )
+            )
+            await _run_blocking(
+                upload_image_build_artifact,
+                create_result.upload,
+                artifact.path,
+                timeout=upload_timeout,
+            )
+            build = await self.complete_image_build(
+                create_result.build.id,
+                CompleteSandboxImageBuildParams(
+                    input_sha256=artifact.sha256_hex,
+                    input_size_bytes=artifact.size_bytes,
+                    input_format=artifact.input_format,
+                ),
+            )
+            if wait:
+                return await self.wait_for_image_build(
+                    build.id,
+                    poll_interval=poll_interval,
+                    timeout=wait_timeout,
+                )
+            return build
+        finally:
+            artifact.cleanup()
+
+    async def build_image_from_dockerfile(
+        self,
+        *,
+        context_path,
+        image_name: str,
+        dockerfile="Dockerfile",
+        docker_tag: Optional[str] = None,
+        platform: str = IMAGE_BUILD_SOURCE_PLATFORM,
+        build_args: Optional[Dict[str, str]] = None,
+        image_init: Optional[SandboxImageInit] = None,
+        image_config_user: Optional[str] = None,
+        wait: bool = True,
+        poll_interval: float = 3.0,
+        wait_timeout: Optional[float] = 35 * 60,
+        temp_dir: Optional[str] = None,
+        upload_timeout: Optional[float] = None,
+    ) -> SandboxImageBuild:
+        tag = docker_tag or make_temp_docker_tag()
+        remove_tag = docker_tag is None
+        try:
+            await _run_blocking(
+                build_docker_image_from_dockerfile,
+                context_path=context_path,
+                dockerfile=dockerfile,
+                tag=tag,
+                platform=platform,
+                build_args=build_args,
+            )
+            return await self.build_image_from_docker_image(
+                docker_image=tag,
+                image_name=image_name,
+                platform=platform,
+                image_init=image_init,
+                image_config_user=image_config_user,
+                wait=wait,
+                poll_interval=poll_interval,
+                wait_timeout=wait_timeout,
+                temp_dir=temp_dir,
+                upload_timeout=upload_timeout,
+            )
+        finally:
+            if remove_tag:
+                await _run_blocking(remove_docker_image, tag)
+
     async def get_runtime_session(self, sandbox_id: str) -> SandboxRuntimeSession:
         detail = await self.get_detail(sandbox_id)
         session = SandboxHandle._to_runtime_session(detail)
@@ -404,6 +625,20 @@ async def create_memory_snapshot(
         )
         return SandboxMemorySnapshotResult(**payload)
 
+    async def update_network(
+        self,
+        sandbox_id: str,
+        policy: SandboxNetworkPolicy,
+    ) -> SandboxNetworkUpdateResult:
+        if not isinstance(policy, SandboxNetworkPolicy):
+            raise TypeError("policy must be a SandboxNetworkPolicy instance")
+        payload = await self._request(
+            "PUT",
+            f"/sandbox/{sandbox_id}/network",
+            data=policy.model_dump(exclude_none=True, by_alias=True),
+        )
+        return SandboxNetworkUpdateResult(**payload)
+
     async def expose(
         self,
         sandbox_id: str,