Skip to content

Commit db0f227

Browse files
sitolesitole
committed
inject access token to request if connection is secured
1 parent b5dfc85 commit db0f227

12 files changed

Lines changed: 158 additions & 31 deletions

File tree

‎packages/python-sdk/e2b/api/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
class SandboxCreateResponse:
2323
sandbox_id: str
2424
envd_version: str
25-
25+
envd_access_token: str
2626

2727
def handle_api_exception(e: Response):
2828
try:

‎packages/python-sdk/e2b/sandbox/sandbox_api.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,10 @@ class SandboxInfo:
2121
"""Sandbox start time."""
2222
end_at: datetime
2323
"""Sandbox expiration date."""
24-
24+
envd_access_token: str
25+
"""EnvD access token."""
26+
envd_version: str
27+
"""EnvD version."""
2528

2629
@dataclass
2730
class SandboxQuery:

‎packages/python-sdk/e2b/sandbox_async/commands/command.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ async def list(
5454
request_timeout=self._connection_config.get_request_timeout(
5555
request_timeout
5656
),
57+
headers=self._base_headers(),
5758
)
5859
return [
5960
ProcessInfo(
@@ -92,6 +93,7 @@ async def kill(
9293
request_timeout=self._connection_config.get_request_timeout(
9394
request_timeout
9495
),
96+
headers=self._base_headers(),
9597
)
9698
return True
9799
except Exception as e:
@@ -124,6 +126,7 @@ async def send_stdin(
124126
request_timeout=self._connection_config.get_request_timeout(
125127
request_timeout
126128
),
129+
headers=self._base_headers(),
127130
)
128131
except Exception as e:
129132
raise handle_rpc_exception(e)
@@ -234,6 +237,7 @@ async def _start(
234237
),
235238
),
236239
headers={
240+
**self._base_headers(),
237241
**authentication_header(user),
238242
KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC),
239243
},
@@ -290,6 +294,7 @@ async def connect(
290294
request_timeout
291295
),
292296
headers={
297+
**self._base_headers(),
293298
KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC),
294299
},
295300
)
@@ -311,3 +316,6 @@ async def connect(
311316
)
312317
except Exception as e:
313318
raise handle_rpc_exception(e)
319+
320+
def _base_headers(self) -> dict:
321+
return self._connection_config.headers

‎packages/python-sdk/e2b/sandbox_async/commands/pty.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ async def kill(
6565
request_timeout=self._connection_config.get_request_timeout(
6666
request_timeout
6767
),
68+
headers=self._base_headers()
6869
)
6970
return True
7071
except Exception as e:
@@ -97,6 +98,7 @@ async def send_stdin(
9798
request_timeout=self._connection_config.get_request_timeout(
9899
request_timeout
99100
),
101+
headers=self._base_headers()
100102
)
101103
except Exception as e:
102104
raise handle_rpc_exception(e)
@@ -139,6 +141,7 @@ async def create(
139141
),
140142
),
141143
headers={
144+
**self._base_headers(),
142145
**authentication_header(user),
143146
KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC),
144147
},
@@ -189,4 +192,8 @@ async def resize(
189192
request_timeout=self._connection_config.get_request_timeout(
190193
request_timeout
191194
),
195+
headers=self._base_headers()
192196
)
197+
198+
def _base_headers(self) -> dict:
199+
return self._connection_config.headers

‎packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py‎

Lines changed: 26 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,7 @@ async def read(
121121
ENVD_API_FILES_ROUTE,
122122
params={"path": path, "username": user},
123123
timeout=self._connection_config.get_request_timeout(request_timeout),
124+
headers=self._base_headers()
124125
)
125126

126127
err = await ahandle_envd_api_exception(r)
@@ -234,6 +235,7 @@ async def write(
234235
files=httpx_files,
235236
params=params,
236237
timeout=self._connection_config.get_request_timeout(request_timeout),
238+
headers=self._base_headers()
237239
)
238240

239241
err = await ahandle_envd_api_exception(r)
@@ -272,7 +274,10 @@ async def list(
272274
request_timeout=self._connection_config.get_request_timeout(
273275
request_timeout
274276
),
275-
headers=authentication_header(user),
277+
headers={
278+
**self._base_headers(),
279+
**authentication_header(user),
280+
}
276281
)
277282

278283
entries: List[EntryInfo] = []
@@ -309,7 +314,10 @@ async def exists(
309314
request_timeout=self._connection_config.get_request_timeout(
310315
request_timeout
311316
),
312-
headers=authentication_header(user),
317+
headers={
318+
**self._base_headers(),
319+
**authentication_header(user),
320+
}
313321
)
314322

315323
return True
@@ -339,7 +347,10 @@ async def remove(
339347
request_timeout=self._connection_config.get_request_timeout(
340348
request_timeout
341349
),
342-
headers=authentication_header(user),
350+
headers={
351+
**self._base_headers(),
352+
**authentication_header(user),
353+
}
343354
)
344355
except Exception as e:
345356
raise handle_rpc_exception(e)
@@ -370,7 +381,10 @@ async def rename(
370381
request_timeout=self._connection_config.get_request_timeout(
371382
request_timeout
372383
),
373-
headers=authentication_header(user),
384+
headers={
385+
**self._base_headers(),
386+
**authentication_header(user),
387+
}
374388
)
375389

376390
return EntryInfo(
@@ -402,7 +416,10 @@ async def make_dir(
402416
request_timeout=self._connection_config.get_request_timeout(
403417
request_timeout
404418
),
405-
headers=authentication_header(user),
419+
headers={
420+
**self._base_headers(),
421+
**authentication_header(user),
422+
}
406423
)
407424

408425
return True
@@ -452,6 +469,7 @@ async def watch_dir(
452469
),
453470
timeout=timeout,
454471
headers={
472+
**self._base_headers(),
455473
**authentication_header(user),
456474
KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC),
457475
},
@@ -468,3 +486,6 @@ async def watch_dir(
468486
return AsyncWatchHandle(events=events, on_event=on_event, on_exit=on_exit)
469487
except Exception as e:
470488
raise handle_rpc_exception(e)
489+
490+
def _base_headers(self) -> dict:
491+
return self._connection_config.headers or {}

‎packages/python-sdk/e2b/sandbox_async/main.py‎

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
from typing import Dict, Optional, TypedDict, overload
55
from typing_extensions import Unpack
66

7+
from e2b.api.client.types import Unset
78
from e2b.connection_config import ConnectionConfig
89
from e2b.envd.api import ENVD_API_HEALTH_ROUTE, ahandle_envd_api_exception
910
from e2b.exceptions import format_request_timeout_error
@@ -32,6 +33,7 @@ async def handle_async_request(self, request):
3233
class AsyncSandboxOpts(TypedDict):
3334
sandbox_id: str
3435
envd_version: Optional[str]
36+
envd_access_token: Optional[str]
3537
connection_config: ConnectionConfig
3638

3739

@@ -105,6 +107,7 @@ def __init__(self, **opts: Unpack[AsyncSandboxOpts]):
105107

106108
self._envd_api_url = f"{'http' if self.connection_config.debug else 'https'}://{self.get_host(self.envd_port)}"
107109
self._envd_version = opts["envd_version"]
110+
self._envd_access_token = opts["envd_access_token"]
108111

109112
self._transport = AsyncTransportWithLogger(limits=self._limits)
110113
self._envd_api = httpx.AsyncClient(
@@ -177,6 +180,7 @@ async def create(
177180
domain: Optional[str] = None,
178181
debug: Optional[bool] = None,
179182
request_timeout: Optional[float] = None,
183+
secure: Optional[bool] = None,
180184
):
181185
"""
182186
Create a new sandbox.
@@ -194,16 +198,12 @@ async def create(
194198
195199
Use this method instead of using the constructor to create a new sandbox.
196200
"""
197-
connection_config = ConnectionConfig(
198-
api_key=api_key,
199-
domain=domain,
200-
debug=debug,
201-
request_timeout=request_timeout,
202-
)
201+
connection_headers = {"x": "b"}
203202

204-
if connection_config.debug:
203+
if debug:
205204
sandbox_id = "debug_sandbox_id"
206205
envd_version = None
206+
envd_access_token = None
207207
else:
208208
response = await SandboxApi._create_sandbox(
209209
template=template or cls.default_template,
@@ -214,13 +214,27 @@ async def create(
214214
debug=debug,
215215
request_timeout=request_timeout,
216216
env_vars=envs,
217+
secure=secure,
217218
)
218219
sandbox_id = response.sandbox_id
219220
envd_version = response.envd_version
221+
envd_access_token = response.envd_access_token
222+
223+
if envd_access_token is not None and not isinstance(envd_access_token, Unset):
224+
connection_headers["X-Access-Token"] = envd_access_token
225+
226+
connection_config = ConnectionConfig(
227+
api_key=api_key,
228+
domain=domain,
229+
debug=debug,
230+
request_timeout=request_timeout,
231+
headers=connection_headers
232+
)
220233

221234
return cls(
222235
sandbox_id=sandbox_id,
223236
envd_version=envd_version,
237+
envd_access_token=envd_access_token,
224238
connection_config=connection_config,
225239
)
226240

@@ -255,9 +269,12 @@ async def connect(
255269
debug=debug,
256270
)
257271

272+
response = await SandboxApi.get_info(sandbox_id)
273+
258274
return cls(
259275
sandbox_id=sandbox_id,
260-
envd_version=None,
276+
envd_version=response.envd_version,
277+
envd_access_token=response.envd_access_token,
261278
connection_config=connection_config,
262279
)
263280

‎packages/python-sdk/e2b/sandbox_async/sandbox_api.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,8 @@ async def list(
8383
),
8484
started_at=sandbox.started_at,
8585
end_at=sandbox.end_at,
86+
envd_access_token=sandbox.envd_access_token,
87+
envd_version=sandbox.envd_version,
8688
)
8789
for sandbox in res.parsed
8890
]
@@ -140,6 +142,8 @@ async def get_info(
140142
),
141143
started_at=res.parsed.started_at,
142144
end_at=res.parsed.end_at,
145+
envd_access_token=res.parsed.envd_access_token,
146+
envd_version=res.parsed.envd_version,
143147
)
144148

145149
@classmethod
@@ -218,6 +222,7 @@ async def _create_sandbox(
218222
timeout: int,
219223
metadata: Optional[Dict[str, str]] = None,
220224
env_vars: Optional[Dict[str, str]] = None,
225+
secure: Optional[bool] = None,
221226
api_key: Optional[str] = None,
222227
domain: Optional[str] = None,
223228
debug: Optional[bool] = None,
@@ -239,6 +244,7 @@ async def _create_sandbox(
239244
metadata=metadata or {},
240245
timeout=timeout,
241246
env_vars=env_vars or {},
247+
secure=secure or False,
242248
),
243249
client=api_client,
244250
)
@@ -267,6 +273,7 @@ async def _create_sandbox(
267273
res.parsed.client_id,
268274
),
269275
envd_version=res.parsed.envd_version,
276+
envd_access_token=res.parsed.envd_access_token,
270277
)
271278

272279
@staticmethod

‎packages/python-sdk/e2b/sandbox_sync/commands/command.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ def list(
5353
request_timeout=self._connection_config.get_request_timeout(
5454
request_timeout
5555
),
56+
headers=self._base_headers(),
5657
)
5758
return [
5859
ProcessInfo(
@@ -91,6 +92,7 @@ def kill(
9192
request_timeout=self._connection_config.get_request_timeout(
9293
request_timeout
9394
),
95+
headers=self._base_headers(),
9496
)
9597
return True
9698
except Exception as e:
@@ -123,6 +125,7 @@ def send_stdin(
123125
request_timeout=self._connection_config.get_request_timeout(
124126
request_timeout
125127
),
128+
headers=self._base_headers(),
126129
)
127130
except Exception as e:
128131
raise handle_rpc_exception(e)
@@ -234,6 +237,7 @@ def _start(
234237
),
235238
),
236239
headers={
240+
**self._base_headers(),
237241
**authentication_header(user),
238242
KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC),
239243
},
@@ -280,6 +284,7 @@ def connect(
280284
process=process_pb2.ProcessSelector(pid=pid),
281285
),
282286
headers={
287+
**self._base_headers(),
283288
KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC),
284289
},
285290
timeout=timeout,
@@ -303,3 +308,6 @@ def connect(
303308
)
304309
except Exception as e:
305310
raise handle_rpc_exception(e)
311+
312+
def _base_headers(self) -> dict:
313+
return self._connection_config.headers

0 commit comments

Comments
 (0)