diff --git a/README.md b/README.md index 081cd3d..71e53aa 100644 --- a/README.md +++ b/README.md @@ -60,6 +60,8 @@ print(f"Verdict: {completed.level}") print(f"MITRE techniques: {completed.mitre_techniques}") ``` +- Pass `auto_select_environment=True` to let the server route the sample to a Windows / Linux / macOS / Android environment based on its detected type. See [recipe 14](./docs/RECIPES.md#14-auto-select-sandbox-environment). + ### Filename canonicalisation (`dynamic_mimetype_check`) By default the server appends the real (Magika-detected) extension to the @@ -158,7 +160,7 @@ for the recommended retry loop. | Document | Purpose | |---|---| -| [Recipes](./docs/RECIPES.md) | 12 runnable examples covering common tasks. | +| [Recipes](./docs/RECIPES.md) | Runnable examples covering common tasks. | | [Type Reference](./docs/TYPES.md) | Pydantic model overview grouped by feature area. | ## Testing utilities diff --git a/docs/RECIPES.md b/docs/RECIPES.md index ca6e8f8..97c8e27 100644 --- a/docs/RECIPES.md +++ b/docs/RECIPES.md @@ -23,6 +23,7 @@ client = ThreatZone(api_key="") 11. [Build a daily report summary](#11-build-a-daily-report-summary) 12. [Cross-reference MITRE techniques with indicators](#12-cross-reference-mitre-techniques-with-indicators) 13. [Control Magika filename canonicalisation per submission (`dynamic_mimetype_check`)](#13-control-dynamic_mimetype_check) +14. [Auto-select sandbox environment](#14-auto-select-sandbox-environment) --- @@ -739,3 +740,32 @@ async with AsyncThreatZone(api_key="...") as client: metafields={"dynamic_mimetype_check": False, "timeout": 120}, ) ``` + +--- + +### 14. Auto-select sandbox environment + +**Goal.** Submit a mixed batch of Windows / Linux / Android samples without hard-coding +a per-sample sandbox OS. Pass `auto_select_environment=True` and let the server route +each file to the right environment based on its detected type. + +```python +from threatzone import ThreatZone + +with ThreatZone(api_key="") as client: + submission = client.create_sandbox_submission( + "/path/to/sample.exe", + auto_select_environment=True, + ) + print(submission.uuid) +``` + +**Notes.** + +- If you also pass `environment="w10_x64"`, the explicit choice wins and the auto flag + is ignored (the server logs a warning but the request succeeds). +- Unknown extensions fall back to your plan's default OS, so the call never fails just + because the file type is unrecognised. +- The same kwarg is accepted by `create_open_in_browser_submission()` for API symmetry; + since there's no file to route on, OIB always falls back to the plan's default + `open_in_browser` OS. diff --git a/examples/recipe_01_submit_and_wait.py b/examples/recipe_01_submit_and_wait.py index 63306b6..4fe2c1c 100644 --- a/examples/recipe_01_submit_and_wait.py +++ b/examples/recipe_01_submit_and_wait.py @@ -30,6 +30,17 @@ def main(client: ThreatZone, sample_path: Path) -> Submission: ) print(f"Created submission {created.uuid}") + # --- Variant: let the server auto-select the environment --- + # When automating across heterogeneous samples (.exe, .elf, .apk, ...), set + # auto_select_environment=True and omit the `environment` kwarg. The server + # routes the sample based on its detected extension. + auto = client.create_sandbox_submission( + sample_path, + auto_select_environment=True, + private=True, + ) + print(f"Submission (auto-env): {auto.uuid}") + try: final = client.wait_for_completion(created.uuid, timeout=900, poll_interval=5) except AnalysisTimeoutError as exc: diff --git a/scripts/sdk-smoke-real-api.py b/scripts/sdk-smoke-real-api.py index c806fb6..bd35525 100644 --- a/scripts/sdk-smoke-real-api.py +++ b/scripts/sdk-smoke-real-api.py @@ -77,6 +77,14 @@ def main() -> int: print(f"Created sandbox: uuid={sandbox_created.uuid} (enqueued)") client.get_submission(sandbox_created.uuid) client.get_network_summary(sandbox_created.uuid) + + # Auto-select environment smoke: verify the server picks an env for us. + auto_submission = client.create_sandbox_submission( + str(file_path), + auto_select_environment=True, + ) + print(f"Auto-select sandbox: uuid={auto_submission.uuid} (enqueued)") + client.get_submission(auto_submission.uuid) except PermissionDeniedError: print("Sandbox submission not allowed by plan (403)") diff --git a/src/threatzone/_async_client.py b/src/threatzone/_async_client.py index 072487e..d29f9f5 100644 --- a/src/threatzone/_async_client.py +++ b/src/threatzone/_async_client.py @@ -187,6 +187,7 @@ async def create_sandbox_submission( file: FileInput, *, environment: str | None = None, + auto_select_environment: bool = False, metafields: dict[str, Any] | list[dict[str, Any]] | None = None, private: bool = False, entrypoint: str | None = None, @@ -199,6 +200,10 @@ async def create_sandbox_submission( Args: file: File to analyze. Can be a file path, bytes, or file-like object. environment: Sandbox environment to use. See get_environments() for options. + Takes precedence over auto_select_environment. + auto_select_environment: When True, the server picks the environment based + on the file's detected extension. Ignored if `environment` is also set. + Falls back to the plan's default OS for unknown extensions. Default: False. metafields: Optional analysis metafields. You can pass either: - a dict map, e.g. {"timeout": 120, "internet_connection": True} @@ -214,6 +219,10 @@ async def create_sandbox_submission( Returns: SubmissionCreated with the new submission UUID. """ + extra_fields: dict[str, Any] = {} + if auto_select_environment: + extra_fields["autoSelectEnvironment"] = auto_select_environment + multipart = self._http._build_multipart_data( file, environment=environment, @@ -222,6 +231,7 @@ async def create_sandbox_submission( entrypoint=entrypoint, password=password, configurations=configurations, + **extra_fields, ) response = await self._http.post("/submissions/sandbox", **multipart) return SubmissionCreated.model_validate(response.json()) @@ -323,6 +333,7 @@ async def create_open_in_browser_submission( url: str, *, environment: str | None = None, + auto_select_environment: bool = False, metafields: dict[str, Any] | list[dict[str, Any]] | None = None, private: bool = False, configurations: dict[str, str] | None = None, @@ -335,6 +346,11 @@ async def create_open_in_browser_submission( Args: url: URL to analyze. environment: Optional OS environment key for browser execution. + Takes precedence over auto_select_environment. + auto_select_environment: When True, server resolves to the plan's + open_in_browser default. (No file extension exists for URL + submissions, so this is effectively API symmetry with sandbox.) + Ignored if `environment` is also set. Default: False. metafields: Optional open_in_browser metafields. You can pass either: - a dict map, e.g. {"timeout": 120} @@ -349,6 +365,8 @@ async def create_open_in_browser_submission( normalized_metafields = _normalize_metafields_json(metafields) if environment is not None: payload["environment"] = environment + if auto_select_environment: + payload["autoSelectEnvironment"] = auto_select_environment if normalized_metafields is not None: payload["metafields"] = normalized_metafields if configurations is not None: diff --git a/src/threatzone/_sync_client.py b/src/threatzone/_sync_client.py index 70f0f47..03cedde 100644 --- a/src/threatzone/_sync_client.py +++ b/src/threatzone/_sync_client.py @@ -184,6 +184,7 @@ def create_sandbox_submission( file: FileInput, *, environment: str | None = None, + auto_select_environment: bool = False, metafields: dict[str, Any] | list[dict[str, Any]] | None = None, private: bool = False, entrypoint: str | None = None, @@ -196,6 +197,10 @@ def create_sandbox_submission( Args: file: File to analyze. Can be a file path, bytes, or file-like object. environment: Sandbox environment to use. See get_environments() for options. + Takes precedence over auto_select_environment. + auto_select_environment: When True, the server picks the environment based + on the file's detected extension. Ignored if `environment` is also set. + Falls back to the plan's default OS for unknown extensions. Default: False. metafields: Optional analysis metafields. You can pass either: - a dict map, e.g. {"timeout": 120, "internet_connection": True} @@ -211,6 +216,10 @@ def create_sandbox_submission( Returns: SubmissionCreated with the new submission UUID. """ + extra_fields: dict[str, Any] = {} + if auto_select_environment: + extra_fields["autoSelectEnvironment"] = auto_select_environment + multipart = self._http._build_multipart_data( file, environment=environment, @@ -219,6 +228,7 @@ def create_sandbox_submission( entrypoint=entrypoint, password=password, configurations=configurations, + **extra_fields, ) response = self._http.post("/submissions/sandbox", **multipart) return SubmissionCreated.model_validate(response.json()) @@ -320,6 +330,7 @@ def create_open_in_browser_submission( url: str, *, environment: str | None = None, + auto_select_environment: bool = False, metafields: dict[str, Any] | list[dict[str, Any]] | None = None, private: bool = False, configurations: dict[str, str] | None = None, @@ -332,6 +343,11 @@ def create_open_in_browser_submission( Args: url: URL to analyze. environment: Optional OS environment key for browser execution. + Takes precedence over auto_select_environment. + auto_select_environment: When True, server resolves to the plan's + open_in_browser default. (No file extension exists for URL + submissions, so this is effectively API symmetry with sandbox.) + Ignored if `environment` is also set. Default: False. metafields: Optional open_in_browser metafields. You can pass either: - a dict map, e.g. {"timeout": 120} @@ -346,6 +362,8 @@ def create_open_in_browser_submission( normalized_metafields = _normalize_metafields_json(metafields) if environment is not None: payload["environment"] = environment + if auto_select_environment: + payload["autoSelectEnvironment"] = auto_select_environment if normalized_metafields is not None: payload["metafields"] = normalized_metafields if configurations is not None: diff --git a/tests/test_async_client.py b/tests/test_async_client.py index 0402c10..33e3fa5 100644 --- a/tests/test_async_client.py +++ b/tests/test_async_client.py @@ -458,6 +458,45 @@ async def test_create_open_in_browser_minimal( assert "environment" not in body assert "metafields" not in body + async def test_create_sandbox_submission_sends_auto_select_environment_when_true( + self, api_key: str, sample_submission_created: dict[str, Any] + ) -> None: + ctx, mock = _patched_async_client(_make_response(200, sample_submission_created)) + with ctx: + async with AsyncThreatZone(api_key=api_key) as client: + await client.create_sandbox_submission( + b"fake-bytes", + auto_select_environment=True, + ) + _, kwargs = mock.call_args + data = kwargs.get("data") or {} + assert data.get("autoSelectEnvironment") == "true" + + async def test_create_sandbox_submission_omits_auto_select_when_false( + self, api_key: str, sample_submission_created: dict[str, Any] + ) -> None: + ctx, mock = _patched_async_client(_make_response(200, sample_submission_created)) + with ctx: + async with AsyncThreatZone(api_key=api_key) as client: + await client.create_sandbox_submission(b"fake-bytes") + _, kwargs = mock.call_args + data = kwargs.get("data") or {} + assert "autoSelectEnvironment" not in data + + async def test_create_open_in_browser_submission_sends_auto_select_environment( + self, api_key: str, sample_submission_created: dict[str, Any] + ) -> None: + ctx, mock = _patched_async_client(_make_response(200, sample_submission_created)) + with ctx: + async with AsyncThreatZone(api_key=api_key) as client: + await client.create_open_in_browser_submission( + "https://example.com", + auto_select_environment=True, + ) + _, kwargs = mock.call_args + body = kwargs["json"] + assert body.get("autoSelectEnvironment") is True + # --------------------------------------------------------------------------- # Submissions - Query diff --git a/tests/test_sync_client.py b/tests/test_sync_client.py index 6c43244..b7a0d83 100644 --- a/tests/test_sync_client.py +++ b/tests/test_sync_client.py @@ -435,6 +435,42 @@ def test_create_open_in_browser_submission_minimal_payload( assert "environment" not in body assert "metafields" not in body + def test_create_sandbox_submission_sends_auto_select_environment_when_true( + self, api_key: str, sample_submission_created: dict[str, Any] + ) -> None: + ctx, mock = _patched_client(_make_response(200, sample_submission_created)) + with ctx, ThreatZone(api_key=api_key) as client: + client.create_sandbox_submission( + b"fake-bytes", + auto_select_environment=True, + ) + _, kwargs = mock.call_args + data = kwargs.get("data") or {} + assert data.get("autoSelectEnvironment") == "true" + + def test_create_sandbox_submission_omits_auto_select_when_false( + self, api_key: str, sample_submission_created: dict[str, Any] + ) -> None: + ctx, mock = _patched_client(_make_response(200, sample_submission_created)) + with ctx, ThreatZone(api_key=api_key) as client: + client.create_sandbox_submission(b"fake-bytes") + _, kwargs = mock.call_args + data = kwargs.get("data") or {} + assert "autoSelectEnvironment" not in data + + def test_create_open_in_browser_submission_sends_auto_select_environment( + self, api_key: str, sample_submission_created: dict[str, Any] + ) -> None: + ctx, mock = _patched_client(_make_response(200, sample_submission_created)) + with ctx, ThreatZone(api_key=api_key) as client: + client.create_open_in_browser_submission( + "https://example.com", + auto_select_environment=True, + ) + _, kwargs = mock.call_args + body = kwargs["json"] + assert body.get("autoSelectEnvironment") is True + # --------------------------------------------------------------------------- # Submissions - Query