diff --git a/.changeset/sandbox-iam-tokens.md b/.changeset/sandbox-iam-tokens.md new file mode 100644 index 0000000000..9883c27160 --- /dev/null +++ b/.changeset/sandbox-iam-tokens.md @@ -0,0 +1,32 @@ +--- +'e2b': minor +'@e2b/python-sdk': minor +--- + +Add the `iam` option to `Sandbox.create` for configuring sandbox workload identity, and a `Secret` class with an `iamToken` / `iam_token` method for defining the workload tokens. Passing a non-empty `tokens` map (name → `{ audience, tokenType }`) enables workload identity for the sandbox: + +```ts +import { Sandbox, Secret } from 'e2b' + +const sandbox = await Sandbox.create({ + iam: { + tokens: { + aws: Secret.iamToken({ audience: 'sts.amazonaws.com', tokenType: 'JWT-SVID' }), + }, + }, +}) +``` + +```python +from e2b import Sandbox, Secret + +sandbox = Sandbox.create( + iam={ + "tokens": { + "aws": Secret.iam_token(audience="sts.amazonaws.com", token_type="JWT-SVID"), + }, + }, +) +``` + +Plain `{ audience, tokenType }` objects (`{"audience": ..., "token_type": ...}` dicts in Python) are accepted as token values too. diff --git a/packages/js-sdk/src/index.ts b/packages/js-sdk/src/index.ts index 074aa3fc9b..54aa9a83ff 100644 --- a/packages/js-sdk/src/index.ts +++ b/packages/js-sdk/src/index.ts @@ -62,6 +62,9 @@ export type { SandboxState, SandboxListOpts, SandboxPaginator, + SandboxIamOpts, + SandboxIamToken, + SandboxIamTokenType, SandboxNetworkOpts, SandboxNetworkInfo, SandboxNetworkSelector, @@ -82,6 +85,8 @@ export type { export type { McpServer } from './sandbox/mcp' +export { Secret } from './secret' + export { ALL_TRAFFIC } from './sandbox/network' export type { diff --git a/packages/js-sdk/src/sandbox/sandboxApi.ts b/packages/js-sdk/src/sandbox/sandboxApi.ts index 8cccdb67d4..a11c955078 100644 --- a/packages/js-sdk/src/sandbox/sandboxApi.ts +++ b/packages/js-sdk/src/sandbox/sandboxApi.ts @@ -200,6 +200,40 @@ export type SandboxNetworkUpdate = { allowInternetAccess?: boolean } +/** + * Workload token type. `'JWT-SVID'` is the only type the API accepts in this + * version; the set is defined server-side and may grow, so any string is + * allowed. + */ +export type SandboxIamTokenType = 'JWT-SVID' | (string & {}) + +/** + * Workload token definition for sandbox workload identity. + */ +export interface SandboxIamToken { + /** + * Audience of the workload token, stored exactly as provided. + */ + audience: string + + /** + * Workload token type. + */ + tokenType: SandboxIamTokenType +} + +/** + * Sandbox workload identity configuration. A non-empty `tokens` map enables + * workload identity for the sandbox. + */ +export interface SandboxIamOpts { + /** + * Named workload-token definitions, keyed by a caller-chosen token name. + * Values can be created with `Secret.iamToken()`. + */ + tokens?: Record +} + /** * What happens when the sandbox timeout is reached. Either the bare action * (`'pause'` / `'kill'`), or an object form that also controls the pause @@ -405,6 +439,26 @@ export interface SandboxOpts extends ConnectionOpts { */ network?: SandboxNetworkOpts + /** + * Sandbox workload identity configuration. Providing a non-empty + * `tokens` map enables workload identity for the sandbox. + * + * @example + * ```ts + * const sandbox = await Sandbox.create({ + * iam: { + * tokens: { + * aws: Secret.iamToken({ + * audience: 'sts.amazonaws.com', + * tokenType: 'JWT-SVID', + * }), + * }, + * }, + * }) + * ``` + */ + iam?: SandboxIamOpts + /** * Volume mounts for the sandbox. * @@ -738,6 +792,41 @@ function buildNetworkBody( } } +function buildIamBody( + iam: SandboxIamOpts | undefined +): components['schemas']['SandboxIam'] | undefined { + // Rebuild the body from the known fields so stray properties on the + // caller's objects never reach the wire and later mutations of the + // caller's map cannot alter the in-flight request. + const tokens: components['schemas']['SandboxIamTokens'] = {} + for (const [name, token] of Object.entries(iam?.tokens ?? {})) { + // Untyped callers can leave holes in the map (conditionally included + // tokens); those don't count toward a non-empty config. + if (!token) { + continue + } + + if ( + typeof token.audience !== 'string' || + typeof token.tokenType !== 'string' + ) { + throw new InvalidArgumentError( + `iam token '${name}' must have string 'audience' and 'tokenType' properties.` + ) + } + + tokens[name] = { audience: token.audience, tokenType: token.tokenType } + } + + // Only a non-empty tokens map enables workload identity, so an empty + // iam config is omitted from the payload entirely. + if (Object.keys(tokens).length === 0) { + return undefined + } + + return { tokens } +} + function buildNetworkUpdateBody( network: SandboxNetworkUpdate ): components['schemas']['SandboxNetworkUpdateConfig'] { @@ -1211,6 +1300,7 @@ export class SandboxApi { secure: opts?.secure ?? true, allow_internet_access: opts?.allowInternetAccess ?? true, network: buildNetworkBody(opts?.network), + iam: buildIamBody(opts?.iam), autoPause: action === 'pause', autoPauseMemory: action === 'pause' ? keepMemory : undefined, autoResume: { enabled: autoResume }, diff --git a/packages/js-sdk/src/secret.ts b/packages/js-sdk/src/secret.ts new file mode 100644 index 0000000000..3347c90c01 --- /dev/null +++ b/packages/js-sdk/src/secret.ts @@ -0,0 +1,32 @@ +import type { SandboxIamToken } from './sandbox/sandboxApi' + +/** + * Secrets and workload identity helpers. + */ +export class Secret { + /** + * Define a workload identity token to pass to `iam.tokens` when creating + * a sandbox. + * + * @param token workload token definition. + * + * @returns a token definition passable to `iam.tokens`. + * + * @example + * ```ts + * const sandbox = await Sandbox.create({ + * iam: { + * tokens: { + * aws: Secret.iamToken({ + * audience: 'sts.amazonaws.com', + * tokenType: 'JWT-SVID', + * }), + * }, + * }, + * }) + * ``` + */ + static iamToken(token: SandboxIamToken): SandboxIamToken { + return { ...token } + } +} diff --git a/packages/js-sdk/tests/sandbox/iam.test.ts b/packages/js-sdk/tests/sandbox/iam.test.ts new file mode 100644 index 0000000000..2cc568d06e --- /dev/null +++ b/packages/js-sdk/tests/sandbox/iam.test.ts @@ -0,0 +1,130 @@ +import { afterAll, afterEach, beforeAll, expect, test } from 'vitest' +import { http, HttpResponse } from 'msw' +import { setupServer } from 'msw/node' + +import { InvalidArgumentError, Sandbox, Secret } from '../../src' +import { TEST_API_KEY, apiUrl } from '../setup' + +let lastCreateBody: Record | undefined + +const server = setupServer( + http.post(apiUrl('/sandboxes'), async ({ request }) => { + lastCreateBody = (await request.json()) as Record + return HttpResponse.json({ + sandboxID: 'test-sandbox-id', + templateID: 'base', + envdVersion: '0.2.4', + }) + }) +) + +beforeAll(() => server.listen({ onUnhandledRequest: 'error' })) + +afterAll(() => server.close()) + +afterEach(() => { + lastCreateBody = undefined + server.resetHandlers() +}) + +test('Sandbox.create sends iam tokens in the request body', async () => { + await Sandbox.create('base', { + apiKey: TEST_API_KEY, + iam: { + tokens: { + aws: { audience: 'sts.amazonaws.com', tokenType: 'JWT-SVID' }, + }, + }, + }) + + expect(lastCreateBody?.iam).toEqual({ + tokens: { + aws: { audience: 'sts.amazonaws.com', tokenType: 'JWT-SVID' }, + }, + }) +}) + +test('Sandbox.create sends Secret.iamToken tokens in the request body', async () => { + await Sandbox.create('base', { + apiKey: TEST_API_KEY, + iam: { + tokens: { + aws: Secret.iamToken({ + audience: 'sts.amazonaws.com', + tokenType: 'JWT-SVID', + }), + }, + }, + }) + + expect(lastCreateBody?.iam).toEqual({ + tokens: { + aws: { audience: 'sts.amazonaws.com', tokenType: 'JWT-SVID' }, + }, + }) +}) + +test('Sandbox.create omits iam from the request body when not provided', async () => { + await Sandbox.create('base', { apiKey: TEST_API_KEY }) + + expect(lastCreateBody).toBeDefined() + expect(lastCreateBody).not.toHaveProperty('iam') +}) + +test('Sandbox.create omits an empty iam config from the request body', async () => { + await Sandbox.create('base', { apiKey: TEST_API_KEY, iam: {} }) + + expect(lastCreateBody).toBeDefined() + expect(lastCreateBody).not.toHaveProperty('iam') + + await Sandbox.create('base', { apiKey: TEST_API_KEY, iam: { tokens: {} } }) + + expect(lastCreateBody).toBeDefined() + expect(lastCreateBody).not.toHaveProperty('iam') +}) + +test('Sandbox.create treats a tokens map with only undefined values as empty', async () => { + await Sandbox.create('base', { + apiKey: TEST_API_KEY, + iam: { + // Untyped callers can build the map conditionally, leaving holes. + tokens: { aws: undefined as never }, + }, + }) + + expect(lastCreateBody).toBeDefined() + expect(lastCreateBody).not.toHaveProperty('iam') +}) + +test('Sandbox.create strips unknown token properties from the request body', async () => { + await Sandbox.create('base', { + apiKey: TEST_API_KEY, + iam: { + tokens: { + aws: { + audience: 'sts.amazonaws.com', + tokenType: 'JWT-SVID', + filePath: '/run/token', + } as never, + }, + }, + }) + + expect(lastCreateBody?.iam).toEqual({ + tokens: { + aws: { audience: 'sts.amazonaws.com', tokenType: 'JWT-SVID' }, + }, + }) +}) + +test('Sandbox.create rejects a token missing audience or tokenType', async () => { + await expect( + Sandbox.create('base', { + apiKey: TEST_API_KEY, + iam: { + // The wire-format casing an untyped caller might copy from a payload. + tokens: { aws: { audience: 'sts.amazonaws.com' } as never }, + }, + }) + ).rejects.toThrowError(InvalidArgumentError) +}) diff --git a/packages/python-sdk/e2b/__init__.py b/packages/python-sdk/e2b/__init__.py index c8b65aa74a..86f912f0e0 100644 --- a/packages/python-sdk/e2b/__init__.py +++ b/packages/python-sdk/e2b/__init__.py @@ -75,6 +75,9 @@ GitHubMcpServer, GitHubMcpServerConfig, McpServer, + SandboxIamOpts, + SandboxIamToken, + SandboxIamTokenType, SandboxInfo, SandboxInfoLifecycle, SandboxMetrics, @@ -98,6 +101,7 @@ from .sandbox_async.main import AsyncSandbox from .sandbox_async.paginator import AsyncSandboxPaginator, AsyncSnapshotPaginator from .sandbox_async.utils import OutputHandler +from .secret import Secret from .sandbox_sync.commands.command_handle import CommandHandle from .sandbox_sync.filesystem.watch_handle import WatchHandle from .sandbox_sync.main import Sandbox @@ -204,6 +208,11 @@ "SandboxLifecycle", "SandboxOnTimeout", "ALL_TRAFFIC", + # IAM + "SandboxIamOpts", + "SandboxIamToken", + "SandboxIamTokenType", + "Secret", # Snapshot "SnapshotInfo", "SnapshotPaginator", diff --git a/packages/python-sdk/e2b/sandbox/sandbox_api.py b/packages/python-sdk/e2b/sandbox/sandbox_api.py index f5aaaffb6a..699ec8040a 100644 --- a/packages/python-sdk/e2b/sandbox/sandbox_api.py +++ b/packages/python-sdk/e2b/sandbox/sandbox_api.py @@ -38,6 +38,15 @@ from e2b.api.client.models import ( SandboxNetworkTransformHeaders as ClientSandboxNetworkTransformHeaders, ) +from e2b.api.client.models import ( + SandboxIam as ClientSandboxIam, +) +from e2b.api.client.models import ( + SandboxIamToken as ClientSandboxIamToken, +) +from e2b.api.client.models import ( + SandboxIamTokens as ClientSandboxIamTokens, +) from e2b.api.client.models import ( SandboxNetworkUpdateConfig, ) @@ -46,6 +55,7 @@ ) from e2b.api.client.types import Unset from e2b.connection_config import ApiParams +from e2b.exceptions import InvalidArgumentException from e2b.sandbox.mcp import McpServer as BaseMcpServer from e2b.sandbox.network import ALL_TRAFFIC from e2b.paginator import PaginatorBase @@ -226,6 +236,38 @@ class SandboxNetworkUpdate(TypedDict, total=False): """ +SandboxIamTokenType = Union[Literal["JWT-SVID"], str] +""" +Workload token type. ``"JWT-SVID"`` is the only type the API accepts in this +version; the set is defined server-side and may grow, so any string is allowed. +""" + + +class SandboxIamToken(TypedDict): + """ + Workload token definition for sandbox workload identity. + """ + + audience: str + """Audience of the workload token, stored exactly as provided.""" + + token_type: SandboxIamTokenType + """Workload token type.""" + + +class SandboxIamOpts(TypedDict, total=False): + """ + Sandbox workload identity configuration. A non-empty ``tokens`` map enables + workload identity for the sandbox. + """ + + tokens: Dict[str, SandboxIamToken] + """ + Named workload-token definitions, keyed by a caller-chosen token name. + Values can be created with :meth:`Secret.iam_token`. + """ + + class SandboxNetworkInfo(TypedDict, total=False): """ Network configuration as returned by the sandbox info endpoint. @@ -406,6 +448,47 @@ def build_network_config( return body +def build_iam_config( + iam: Optional[SandboxIamOpts], +) -> Optional[ClientSandboxIam]: + """Resolve a :class:`SandboxIamOpts` into the API client body. + + Returns ``None`` for a config with no tokens — only a non-empty tokens + map enables workload identity, so an empty config is omitted from the + request payload entirely. + """ + if iam is None: + return None + + tokens = iam.get("tokens") + if not tokens: + return None + + client_tokens = ClientSandboxIamTokens() + for name, token in tokens.items(): + # Re-check at runtime for callers that bypass the TypedDict — a bare + # KeyError from deep inside create would not name the option or the + # snake_case key the wire's camelCase 'tokenType' maps to. + if ( + not isinstance(token, dict) + or "audience" not in token + or "token_type" not in token + ): + raise InvalidArgumentException( + f"iam token {name!r} must be a dict with 'audience' and " + "'token_type' keys (snake_case 'token_type', not 'tokenType')." + ) + + client_tokens[name] = ClientSandboxIamToken( + audience=token["audience"], + token_type=token["token_type"], + ) + + body = ClientSandboxIam() + body.tokens = client_tokens + return body + + def build_network_update_body( network: SandboxNetworkUpdate, ) -> SandboxNetworkUpdateConfig: diff --git a/packages/python-sdk/e2b/sandbox_async/main.py b/packages/python-sdk/e2b/sandbox_async/main.py index 40ff5a2e1f..f22eaa729d 100644 --- a/packages/python-sdk/e2b/sandbox_async/main.py +++ b/packages/python-sdk/e2b/sandbox_async/main.py @@ -22,6 +22,7 @@ from e2b.sandbox.main import SandboxOpts from e2b.sandbox.sandbox_api import ( McpServer, + SandboxIamOpts, SandboxLifecycle, SandboxMetrics, SandboxNetworkOpts, @@ -179,6 +180,7 @@ async def create( allow_internet_access: bool = True, mcp: Optional[McpServer] = None, network: Optional[SandboxNetworkOpts] = None, + iam: Optional[SandboxIamOpts] = None, lifecycle: Optional[SandboxLifecycle] = None, volume_mounts: Optional[SandboxAsyncVolumeMount] = None, logger: Optional[logging.Logger] = None, @@ -197,6 +199,7 @@ async def create( :param allow_internet_access: Allow sandbox to access the internet, defaults to `True`. If set to `False`, it works the same as setting network `deny_out` to `[0.0.0.0/0]`. :param mcp: MCP server to enable in the sandbox :param network: Sandbox network configuration. ``allow_out``/``deny_out`` may also be a callable receiving a :class:`SandboxNetworkSelectorContext` (``ctx.all_traffic``, ``ctx.rules``) and returning a list of strings. Per-host transform rules are nested under ``network.rules``. + :param iam: Sandbox workload identity configuration. A non-empty ``tokens`` map enables workload identity for the sandbox; token definitions can be created with :meth:`Secret.iam_token`. Example: ``{"tokens": {"aws": Secret.iam_token(audience="sts.amazonaws.com", token_type="JWT-SVID")}}`` :param lifecycle: Sandbox lifecycle configuration — ``on_timeout``: ``"kill"`` (default) or ``"pause"``, or an object ``{"action": "pause"|"kill", "keep_memory": bool}`` where ``keep_memory`` (default ``True``) set to ``False`` makes a timeout auto-pause filesystem-only (cold-boots on resume; cannot be combined with ``auto_resume``); ``auto_resume``: ``False`` (default) or ``True`` (only when ``on_timeout`` action is ``"pause"``). Example: ``{"on_timeout": {"action": "pause", "keep_memory": False}}`` :param volume_mounts: Dictionary mapping mount paths to AsyncVolume instances or volume names :param logger: Logger used for request and response logging for this sandbox. Accepts any standard library `logging.Logger`. When omitted, no request/response logging is emitted. @@ -229,6 +232,7 @@ async def create( allow_internet_access=allow_internet_access, mcp=mcp, network=network, + iam=iam, lifecycle=lifecycle, volume_mounts=transformed_mounts, logger=logger, @@ -1099,6 +1103,7 @@ async def _create( allow_internet_access: bool, mcp: Optional[McpServer] = None, network: Optional[SandboxNetworkOpts] = None, + iam: Optional[SandboxIamOpts] = None, lifecycle: Optional[SandboxLifecycle] = None, volume_mounts: Optional[list] = None, logger: Optional[logging.Logger] = None, @@ -1123,6 +1128,7 @@ async def _create( allow_internet_access=allow_internet_access, mcp=mcp, network=network, + iam=iam, lifecycle=lifecycle, volume_mounts=volume_mounts, logger=logger, diff --git a/packages/python-sdk/e2b/sandbox_async/sandbox_api.py b/packages/python-sdk/e2b/sandbox_async/sandbox_api.py index 9df7a5a956..a540e8c523 100644 --- a/packages/python-sdk/e2b/sandbox_async/sandbox_api.py +++ b/packages/python-sdk/e2b/sandbox_async/sandbox_api.py @@ -49,6 +49,7 @@ from e2b.sandbox.sandbox_api import ( build_network_update_body, McpServer, + SandboxIamOpts, SandboxInfo, SandboxLifecycle, SandboxMetrics, @@ -56,6 +57,7 @@ SandboxNetworkUpdate, SandboxQuery, SnapshotInfo, + build_iam_config, build_network_config, ) from e2b.sandbox_async.paginator import AsyncSandboxPaginator @@ -207,6 +209,7 @@ async def _create_sandbox( secure: bool, mcp: Optional[McpServer] = None, network: Optional[SandboxNetworkOpts] = None, + iam: Optional[SandboxIamOpts] = None, lifecycle: Optional[SandboxLifecycle] = None, volume_mounts: Optional[List[SandboxVolumeMountAPI]] = None, logger: Optional[logging.Logger] = None, @@ -267,6 +270,7 @@ async def _create_sandbox( secure=secure, allow_internet_access=allow_internet_access, network=SandboxNetworkConfig(**network_body) if network_body else UNSET, + iam=build_iam_config(iam) or UNSET, volume_mounts=volume_mounts if volume_mounts else UNSET, ) diff --git a/packages/python-sdk/e2b/sandbox_sync/main.py b/packages/python-sdk/e2b/sandbox_sync/main.py index 4141eeda86..db4d46d3ba 100644 --- a/packages/python-sdk/e2b/sandbox_sync/main.py +++ b/packages/python-sdk/e2b/sandbox_sync/main.py @@ -20,6 +20,7 @@ from e2b.sandbox.main import SandboxOpts from e2b.sandbox.sandbox_api import ( McpServer, + SandboxIamOpts, SandboxLifecycle, SandboxMetrics, SandboxNetworkOpts, @@ -168,6 +169,7 @@ def create( allow_internet_access: bool = True, mcp: Optional[McpServer] = None, network: Optional[SandboxNetworkOpts] = None, + iam: Optional[SandboxIamOpts] = None, lifecycle: Optional[SandboxLifecycle] = None, volume_mounts: Optional[SandboxVolumeMount] = None, logger: Optional[logging.Logger] = None, @@ -186,6 +188,7 @@ def create( :param allow_internet_access: Allow sandbox to access the internet, defaults to `True`. If set to `False`, it works the same as setting network `deny_out` to `[0.0.0.0/0]`. :param mcp: MCP server to enable in the sandbox :param network: Sandbox network configuration. ``allow_out``/``deny_out`` may also be a callable receiving a :class:`SandboxNetworkSelectorContext` (``ctx.all_traffic``, ``ctx.rules``) and returning a list of strings. Per-host transform rules are nested under ``network.rules``. + :param iam: Sandbox workload identity configuration. A non-empty ``tokens`` map enables workload identity for the sandbox; token definitions can be created with :meth:`Secret.iam_token`. Example: ``{"tokens": {"aws": Secret.iam_token(audience="sts.amazonaws.com", token_type="JWT-SVID")}}`` :param lifecycle: Sandbox lifecycle configuration — ``on_timeout``: ``"kill"`` (default) or ``"pause"``, or an object ``{"action": "pause"|"kill", "keep_memory": bool}`` where ``keep_memory`` (default ``True``) set to ``False`` makes a timeout auto-pause filesystem-only (cold-boots on resume; cannot be combined with ``auto_resume``); ``auto_resume``: ``False`` (default) or ``True`` (only when ``on_timeout`` action is ``"pause"``). Example: ``{"on_timeout": {"action": "pause", "keep_memory": False}}`` :param volume_mounts: Dictionary mapping mount paths to Volume instances or volume names :param logger: Logger used for request and response logging for this sandbox. Accepts any standard library `logging.Logger`. When omitted, no request/response logging is emitted. @@ -218,6 +221,7 @@ def create( allow_internet_access=allow_internet_access, mcp=mcp, network=network, + iam=iam, lifecycle=lifecycle, volume_mounts=transformed_mounts, logger=logger, @@ -1090,6 +1094,7 @@ def _create( allow_internet_access: bool, mcp: Optional[McpServer] = None, network: Optional[SandboxNetworkOpts] = None, + iam: Optional[SandboxIamOpts] = None, lifecycle: Optional[SandboxLifecycle] = None, volume_mounts: Optional[list] = None, logger: Optional[logging.Logger] = None, @@ -1114,6 +1119,7 @@ def _create( allow_internet_access=allow_internet_access, mcp=mcp, network=network, + iam=iam, lifecycle=lifecycle, volume_mounts=volume_mounts, logger=logger, diff --git a/packages/python-sdk/e2b/sandbox_sync/sandbox_api.py b/packages/python-sdk/e2b/sandbox_sync/sandbox_api.py index c7bf34e827..d11876f321 100644 --- a/packages/python-sdk/e2b/sandbox_sync/sandbox_api.py +++ b/packages/python-sdk/e2b/sandbox_sync/sandbox_api.py @@ -48,6 +48,7 @@ from e2b.sandbox.sandbox_api import ( build_network_update_body, McpServer, + SandboxIamOpts, SandboxInfo, SandboxLifecycle, SandboxMetrics, @@ -55,6 +56,7 @@ SandboxNetworkUpdate, SandboxQuery, SnapshotInfo, + build_iam_config, build_network_config, ) from e2b.sandbox_sync.paginator import SandboxPaginator, get_api_client @@ -206,6 +208,7 @@ def _create_sandbox( secure: bool, mcp: Optional[McpServer] = None, network: Optional[SandboxNetworkOpts] = None, + iam: Optional[SandboxIamOpts] = None, lifecycle: Optional[SandboxLifecycle] = None, volume_mounts: Optional[List[SandboxVolumeMountAPI]] = None, logger: Optional[logging.Logger] = None, @@ -266,6 +269,7 @@ def _create_sandbox( secure=secure, allow_internet_access=allow_internet_access, network=SandboxNetworkConfig(**network_body) if network_body else UNSET, + iam=build_iam_config(iam) or UNSET, volume_mounts=volume_mounts if volume_mounts else UNSET, ) diff --git a/packages/python-sdk/e2b/secret.py b/packages/python-sdk/e2b/secret.py new file mode 100644 index 0000000000..d895bc468c --- /dev/null +++ b/packages/python-sdk/e2b/secret.py @@ -0,0 +1,34 @@ +from e2b.sandbox.sandbox_api import SandboxIamToken, SandboxIamTokenType + + +class Secret: + """ + Secrets and workload identity helpers. + """ + + @staticmethod + def iam_token(*, audience: str, token_type: SandboxIamTokenType) -> SandboxIamToken: + """ + Define a workload identity token to pass to ``iam.tokens`` when + creating a sandbox. + + :param audience: Audience of the workload token, stored exactly as provided. + :param token_type: Workload token type. + + :return: A token definition passable to ``iam.tokens``. + + Example: + ```python + sandbox = Sandbox.create( + iam={ + "tokens": { + "aws": Secret.iam_token( + audience="sts.amazonaws.com", + token_type="JWT-SVID", + ), + }, + }, + ) + ``` + """ + return SandboxIamToken(audience=audience, token_type=token_type) diff --git a/packages/python-sdk/tests/async/sandbox_async/test_create.py b/packages/python-sdk/tests/async/sandbox_async/test_create.py index aa17c2c52f..3b18c6f884 100644 --- a/packages/python-sdk/tests/async/sandbox_async/test_create.py +++ b/packages/python-sdk/tests/async/sandbox_async/test_create.py @@ -4,12 +4,14 @@ import httpx import pytest -from e2b import AsyncSandbox, SandboxQuery, SandboxState +from e2b import AsyncSandbox, SandboxQuery, SandboxState, Secret from e2b.api.client.models import ( NewSandbox, SandboxAutoResumeConfig, ) +from e2b.api.client.types import UNSET from e2b.exceptions import InvalidArgumentException +from e2b.sandbox.sandbox_api import build_iam_config @pytest.mark.skip_debug() @@ -60,6 +62,78 @@ def test_create_payload_deserializes_auto_resume_enabled(): assert body.auto_resume.to_dict() == {"enabled": False} +def test_create_payload_serializes_iam_tokens(): + iam = build_iam_config( + { + "tokens": { + "aws": {"audience": "sts.amazonaws.com", "token_type": "JWT-SVID"}, + }, + } + ) + assert iam is not None + + body = NewSandbox(template_id="template-id", iam=iam) + + assert body.to_dict()["iam"] == { + "tokens": { + "aws": {"audience": "sts.amazonaws.com", "tokenType": "JWT-SVID"}, + }, + } + + +def test_create_payload_serializes_secret_iam_token(): + iam = build_iam_config( + { + "tokens": { + "aws": Secret.iam_token( + audience="sts.amazonaws.com", token_type="JWT-SVID" + ), + }, + } + ) + assert iam is not None + + body = NewSandbox(template_id="template-id", iam=iam) + + assert body.to_dict()["iam"] == { + "tokens": { + "aws": {"audience": "sts.amazonaws.com", "tokenType": "JWT-SVID"}, + }, + } + + +def test_create_payload_omits_iam_when_not_provided_or_empty(): + assert build_iam_config(None) is None + assert build_iam_config({}) is None + assert build_iam_config({"tokens": {}}) is None + + body = NewSandbox(template_id="template-id", iam=UNSET) + + assert "iam" not in body.to_dict() + + +def test_create_payload_rejects_malformed_iam_tokens(): + # The wire-format casing a user might copy from the JS example or a + # serialized payload must fail with an actionable error, not a KeyError. + with pytest.raises(InvalidArgumentException, match="token_type"): + build_iam_config( + cast( + Any, + { + "tokens": { + "aws": { + "audience": "sts.amazonaws.com", + "tokenType": "JWT-SVID", + }, + }, + }, + ) + ) + + with pytest.raises(InvalidArgumentException): + build_iam_config(cast(Any, {"tokens": {"aws": None}})) + + @pytest.mark.skip_debug() async def test_filesystem_only_auto_pause_rejects_auto_resume(): # A filesystem-only auto-pause snapshot can only be resumed explicitly, so diff --git a/packages/python-sdk/tests/sync/sandbox_sync/test_create.py b/packages/python-sdk/tests/sync/sandbox_sync/test_create.py index 737e346d1a..e3c2ba2fa6 100644 --- a/packages/python-sdk/tests/sync/sandbox_sync/test_create.py +++ b/packages/python-sdk/tests/sync/sandbox_sync/test_create.py @@ -4,13 +4,14 @@ import httpx import pytest -from e2b import Sandbox, SandboxState +from e2b import Sandbox, SandboxState, Secret from e2b.api.client.models import ( NewSandbox, SandboxAutoResumeConfig, ) +from e2b.api.client.types import UNSET from e2b.exceptions import InvalidArgumentException -from e2b.sandbox.sandbox_api import SandboxQuery +from e2b.sandbox.sandbox_api import SandboxQuery, build_iam_config @pytest.mark.skip_debug() @@ -61,6 +62,78 @@ def test_create_payload_deserializes_auto_resume_enabled(): assert body.auto_resume.to_dict() == {"enabled": False} +def test_create_payload_serializes_iam_tokens(): + iam = build_iam_config( + { + "tokens": { + "aws": {"audience": "sts.amazonaws.com", "token_type": "JWT-SVID"}, + }, + } + ) + assert iam is not None + + body = NewSandbox(template_id="template-id", iam=iam) + + assert body.to_dict()["iam"] == { + "tokens": { + "aws": {"audience": "sts.amazonaws.com", "tokenType": "JWT-SVID"}, + }, + } + + +def test_create_payload_serializes_secret_iam_token(): + iam = build_iam_config( + { + "tokens": { + "aws": Secret.iam_token( + audience="sts.amazonaws.com", token_type="JWT-SVID" + ), + }, + } + ) + assert iam is not None + + body = NewSandbox(template_id="template-id", iam=iam) + + assert body.to_dict()["iam"] == { + "tokens": { + "aws": {"audience": "sts.amazonaws.com", "tokenType": "JWT-SVID"}, + }, + } + + +def test_create_payload_omits_iam_when_not_provided_or_empty(): + assert build_iam_config(None) is None + assert build_iam_config({}) is None + assert build_iam_config({"tokens": {}}) is None + + body = NewSandbox(template_id="template-id", iam=UNSET) + + assert "iam" not in body.to_dict() + + +def test_create_payload_rejects_malformed_iam_tokens(): + # The wire-format casing a user might copy from the JS example or a + # serialized payload must fail with an actionable error, not a KeyError. + with pytest.raises(InvalidArgumentException, match="token_type"): + build_iam_config( + cast( + Any, + { + "tokens": { + "aws": { + "audience": "sts.amazonaws.com", + "tokenType": "JWT-SVID", + }, + }, + }, + ) + ) + + with pytest.raises(InvalidArgumentException): + build_iam_config(cast(Any, {"tokens": {"aws": None}})) + + @pytest.mark.skip_debug() def test_filesystem_only_auto_pause_rejects_auto_resume(): # A filesystem-only auto-pause snapshot can only be resumed explicitly, so