"""Unit tests for the SDK-backed ThreatZone integration. The integration delegates all HTTP work to the official `threatzone` Python SDK, so these tests mock the SDK client methods returning pydantic models straight from the SDK's own type system. """ from __future__ import annotations import hashlib import unittest from pathlib import Path from unittest.mock import MagicMock, call, mock_open, patch import pytest import ThreatZone as integration from CommonServerPython import DemistoException from threatzone import ( AnalysisTimeoutError, APIError, AuthenticationError, NotFoundError, ReportUnavailableError, YaraRulePendingError, ) from threatzone import ThreatZone as ThreatZoneSDK from threatzone.testing import FakeThreatZoneAPI, scenarios from threatzone.types import ( Artifact, ArtifactsResponse, ExtractedConfigsResponse, Indicator, IndicatorsResponse, IoC, IoCsResponse, SubmissionCreated, UserInfo, YaraRule, YaraRulesResponse, ) from threatzone.types.indicators import ArtifactHashes, IndicatorLevels from threatzone.types.config import MetafieldOption def _user_info(api_used: int = 5, daily_used: int = 5, concurrent_used: int = 1) -> UserInfo: return UserInfo.model_validate( { "userInfo": { "email": "name@company.com", "fullName": "Test User", "workspace": { "id": "ws-1", "name": "ACME Lab", "alias": "acme", "private": True, "type": "organization", }, "limitsCount": { "apiRequestCount": api_used, "dailySubmissionCount": daily_used, "concurrentSubmissionCount": concurrent_used, }, }, "plan": { "planName": "Enterprise", "startTime": "2025-01-01", "endTime": "2026-01-01", "subsTime": "yearly", "fileLimits": {"extensions": ["exe", "dll"], "fileSize": "256 MiB"}, "submissionLimits": { "apiLimit": 9999, "dailyLimit": 999, "concurrentLimit": 2, }, }, "modules": [ { "moduleId": "m1", "moduleName": "Sandbox", "startTime": "2025-01-01", "endTime": "2026-01-01", }, { "moduleId": "m2", "moduleName": "CDR", "startTime": "2025-01-01", "endTime": "2026-01-01", }, ], } ) def _submission(level: str = "malicious", status: str = "completed", report_type: str = "dynamic"): """Build a minimal Submission via SDK model validation.""" from threatzone.types.submissions import Submission payload = { "uuid": "c89d310b-7862-4534-998a-3eb39d9a9d42", "type": "file", "filename": "sample.exe", "hashes": { "md5": "5d41402abc4b2a76b9719d911017c592", "sha1": "aaf4c61ddcc5e8a2dabede0f3b482cd9aea9434d", "sha256": "6e899ff7ef160d96787f505b7a9e17b789695bf3206a130c462b3820898257d0", }, "level": level, "private": True, "tags": [], "reports": [{"type": report_type, "status": status, "level": level}], "overview": {"status": "completed"}, "indicators": { "levels": {"malicious": 1, "suspicious": 0, "benign": 0}, "artifactCount": 0, }, "mitreTechniques": [], "createdAt": "2025-01-01T00:00:00Z", "updatedAt": "2025-01-01T00:00:00Z", } return Submission.model_validate(payload) def _make_client() -> integration.Client: """Build a Client without spinning up the real httpx session.""" with ( patch.object(integration, "httpx") as mock_httpx, patch.object(integration, "ThreatZoneSDK") as mock_sdk, ): mock_httpx.Client.return_value = MagicMock() mock_sdk.return_value = MagicMock() return integration.Client(base_url="https://app.threat.zone", api_key="key", verify=True, proxy=False) def _metafield( key: str, default: bool | int | str, *, active: bool = True, accessible: bool = True, ) -> MetafieldOption: return MetafieldOption.model_validate( { "key": key, "label": key, "description": key, "type": "select", "default": default, "active": active, "accessible": accessible, "options": None, } ) # --------------------------------------------------------------------------- # Helper unit tests # --------------------------------------------------------------------------- class TestPureHelpers(unittest.TestCase): def test_normalize_sdk_base_url(self): assert integration.normalize_sdk_base_url("https://app.threat.zone") == "https://app.threat.zone/public-api" assert integration.normalize_sdk_base_url("https://app.threat.zone/public-api/") == "https://app.threat.zone/public-api" def test_extract_api_key_supports_type_9_and_legacy_values(self): assert integration.extract_api_key({"identifier": "", "password": "type-9-key"}) == "type-9-key" assert integration.extract_api_key("legacy-key") == "legacy-key" assert integration.extract_api_key(None) == "" def test_normalize_sdk_base_url_rejects_missing_url(self): with pytest.raises(DemistoException, match="Server URL is required"): integration.normalize_sdk_base_url(" ") def test_parse_json_object_argument(self): assert integration.parse_json_object_argument('{"networkConfig":"id"}', "configurations") == {"networkConfig": "id"} with pytest.raises(DemistoException, match="JSON object"): integration.parse_json_object_argument("[]", "configurations") def test_parse_csv_and_bounded_integer_arguments(self): assert integration.parse_csv_list_argument(" malicious, suspicious, ") == [ "malicious", "suspicious", ] assert integration.parse_bounded_int_argument("500", "limit", minimum=1, maximum=500) == 500 with pytest.raises(DemistoException, match="between 1 and 500"): integration.parse_bounded_int_argument("501", "limit", minimum=1, maximum=500) def test_report_pagination_has_a_hard_page_limit(self): fetch_page = MagicMock() fetch_page.return_value = MagicMock( items=[object()], total=integration.MAX_REPORT_FINDINGS_PAGES * integration.REPORT_FINDINGS_PAGE_SIZE + 1, ) with pytest.raises(DemistoException, match="pagination exceeded"): integration.get_all_report_items(fetch_page, "submission-uuid") assert fetch_page.call_count == integration.MAX_REPORT_FINDINGS_PAGES def test_translate_score_levels(self): assert integration.translate_score(None) == 0 assert integration.translate_score(0) == 0 assert integration.translate_score(1) == 1 assert integration.translate_score(2) == 2 assert integration.translate_score(3) == 3 assert integration.translate_score(99) == 3 def test_get_reputation_reliability_known(self): assert integration.get_reputation_reliability("A - Completely reliable") == "A - Completely reliable" def test_get_reputation_reliability_unknown(self): assert integration.get_reputation_reliability("not-a-real-value") is None def test_parse_modules_argument_csv(self): assert integration.parse_modules_argument("a, b, ,c") == ["a", "b", "c"] def test_parse_modules_argument_json_array(self): assert integration.parse_modules_argument('["a","b"]') == ["a", "b"] def test_parse_modules_argument_empty(self): assert integration.parse_modules_argument(None) is None assert integration.parse_modules_argument("") is None def test_parse_analyze_config_valid(self): result = integration.parse_analyze_config_argument('[{"metafieldId":"x","value":1}]') assert result == [{"metafieldId": "x", "value": 1}] def test_parse_analyze_config_invalid_json(self): with pytest.raises(DemistoException): integration.parse_analyze_config_argument("not-json") def test_metafields_from_legacy_args_merges(self): fields = integration.metafields_from_legacy_args( { "timeout": "120", "work_path": "desktop", "mouse_simulation": "true", "raw_logs": "true", "modules": '["cdr"]', "analyze_config": '[{"metafieldId":"timeout","value":300}]', } ) assert fields["timeout"] == 300 # user override wins assert fields["work_path"] == "desktop" assert fields["mouse_simulation"] is True assert "raw_logs" not in fields assert "modules" not in fields def test_metafields_from_legacy_args_uses_api_defaults(self): assert integration.metafields_from_legacy_args({}, {"snapshot": True}) == {"snapshot": True} def test_metafields_from_legacy_args_forwards_explicit_false(self): assert integration.metafields_from_legacy_args({"snapshot": "false"}, {"snapshot": True}) == {"snapshot": False} def test_sandbox_api_defaults_filters_unavailable_definitions(self): client = _make_client() client.sdk.get_metafields.return_value = [ _metafield("snapshot", True), _metafield("inactive", True, active=False), _metafield("inaccessible", True, accessible=False), ] assert integration.sandbox_api_defaults(client) == {"snapshot": True} client.sdk.get_metafields.assert_called_once_with("sandbox") def test_submission_level_int(self): assert integration.submission_level_int("malicious") == 3 assert integration.submission_level_int("suspicious") == 2 assert integration.submission_level_int("benign") == 1 assert integration.submission_level_int("unknown") == 0 assert integration.submission_level_int(None) is None def test_report_status_int(self): assert integration.report_status_int("completed") == 5 assert integration.report_status_int("in_progress") == 3 assert integration.report_status_int(None) is None def test_parse_file_size_mib(self): assert integration.parse_file_size_mib("256 MiB") == 256 assert integration.parse_file_size_mib(128) == 128 assert integration.parse_file_size_mib("1.5 MiB") == 1.5 assert integration.parse_file_size_mib("unknown") is None class TestClient(unittest.TestCase): @patch.object(integration, "ThreatZoneSDK") @patch.object(integration.httpx, "Client") def test_supplied_http_client_matches_sdk_transport_defaults(self, http_client_mock, sdk_mock): client = integration.Client( base_url="https://app.threat.zone", api_key="key", verify=False, proxy=True, reliability="A - Completely reliable", ) assert isinstance(client, integration.ContentClient) assert client.reliability == "A - Completely reliable" http_client_kwargs = http_client_mock.call_args.kwargs assert http_client_kwargs["verify"] is False assert http_client_kwargs["trust_env"] is True assert http_client_kwargs["follow_redirects"] is True timeout = http_client_kwargs["timeout"] assert timeout.connect == integration.SDK_REQUEST_TIMEOUT_SECONDS assert timeout.read == integration.SDK_REQUEST_TIMEOUT_SECONDS assert timeout.write == integration.SDK_REQUEST_TIMEOUT_SECONDS assert timeout.pool == integration.SDK_REQUEST_TIMEOUT_SECONDS assert sdk_mock.call_args.kwargs["http_client"] is http_client_mock.return_value client.close() # --------------------------------------------------------------------------- # Command-handler tests # --------------------------------------------------------------------------- class TestCheckLimits(unittest.TestCase): def setUp(self): self.client = _make_client() def test_basic(self): self.client.sdk.get_user_info.return_value = _user_info() results = integration.threatzone_check_limits(self.client, {}) assert len(results) == 1 outputs = results[0].outputs assert outputs["E_Mail"] == "name@company.com" assert outputs["API_Limit"] == "5/9999" def test_detailed_adds_plan_and_metadata(self): self.client.sdk.get_user_info.return_value = _user_info() results = integration.threatzone_check_limits(self.client, {"detailed": "true"}) prefixes = [r.outputs_prefix for r in results] assert "ThreatZone.Plan" in prefixes assert "ThreatZone.Metadata" in prefixes plan_result = next(result for result in results if result.outputs_prefix == "ThreatZone.Plan") assert plan_result.outputs["File_Size_Limit_MiB"] == 256 class TestPlanCapacity(unittest.TestCase): def setUp(self): self.client = _make_client() def test_concurrent_limit_blocks_sandbox(self): self.client.sdk.get_user_info.return_value = _user_info(concurrent_used=2) with pytest.raises(DemistoException) as excinfo: integration._verify_plan_capacity(self.client, requires_concurrent=True) assert "Concurrent" in str(excinfo.value) def test_concurrent_limit_does_not_block_url(self): self.client.sdk.get_user_info.return_value = _user_info(concurrent_used=2) integration._verify_plan_capacity(self.client, requires_concurrent=False) def test_api_limit_blocks(self): self.client.sdk.get_user_info.return_value = _user_info(api_used=9999) with pytest.raises(DemistoException) as excinfo: integration._verify_plan_capacity(self.client, requires_concurrent=False) assert "API request limit" in str(excinfo.value) class TestUrlSubmission(unittest.TestCase): def setUp(self): self.client = _make_client() self.client.sdk.get_user_info.return_value = _user_info() def test_url_submission_returns_uuid(self): self.client.sdk.create_url_submission.return_value = SubmissionCreated(uuid="abc", message="ok") results = integration.threatzone_submit_url_analysis(self.client, {"url": "https://example.com", "private": "true"}) primary = results[0] assert primary.outputs["UUID"] == "abc" assert primary.outputs["URL"] == "https://example.com" self.client.sdk.create_url_submission.assert_called_once_with("https://example.com", private=True, safe_browsing=False) def test_url_submission_forwards_safe_browsing(self): self.client.sdk.create_url_submission.return_value = SubmissionCreated(uuid="abc", message="ok") integration.threatzone_submit_url_analysis( self.client, {"url": "https://example.com", "safe_browsing": "true"}, ) assert self.client.sdk.create_url_submission.call_args.kwargs["safe_browsing"] is True def test_url_submission_requires_url(self): with pytest.raises(DemistoException): integration.threatzone_submit_url_analysis(self.client, {}) class TestSandboxUpload(unittest.TestCase): def setUp(self): self.client = _make_client() self.client.sdk.get_user_info.return_value = _user_info() self.client.sdk.get_metafields.return_value = [] self.client.sdk.create_sandbox_submission.return_value = SubmissionCreated(uuid="sb-uuid", message="ok") @patch.object(integration, "demisto") def test_sandbox_upload(self, mock_demisto): mock_demisto.getFilePath.return_value = { "path": "/tmp/generated-entry-id", "name": "original sample.exe", } with patch.object(integration.Path, "open", mock_open(read_data=b"sample")): results = integration.threatzone_sandbox_upload_sample( self.client, { "entry_id": "1", "environment": "w10_x64", "private": "true", "timeout": "120", "configurations": '{"startArguments":"--safe"}', }, ) assert results[0].outputs["UUID"] == "sb-uuid" assert results[0].outputs["FileName"] == "original sample.exe" upload_file = self.client.sdk.create_sandbox_submission.call_args.args[0] assert upload_file.name == "original sample.exe" call_kwargs = self.client.sdk.create_sandbox_submission.call_args.kwargs assert call_kwargs["environment"] == "w10_x64" assert call_kwargs["auto_select_environment"] is False assert call_kwargs["private"] is True assert call_kwargs["metafields"]["timeout"] == 120 assert call_kwargs["configurations"] == {"startArguments": "--safe"} @patch.object(integration, "demisto") def test_sandbox_auto_environment_ignores_explicit_default(self, mock_demisto): mock_demisto.getFilePath.return_value = { "path": "/tmp/sample.exe", "name": "sample.exe", } with patch.object(integration.Path, "open", mock_open(read_data=b"sample")): integration.threatzone_sandbox_upload_sample( self.client, {"entry_id": "1", "environment": "w7_x64", "auto": "true"}, ) call_kwargs = self.client.sdk.create_sandbox_submission.call_args.kwargs assert call_kwargs["environment"] is None assert call_kwargs["auto_select_environment"] is True @patch.object(integration, "demisto") def test_bat_sandbox_uses_sdk_with_api_defaults(self, mock_demisto): mock_demisto.getFilePath.return_value = { "path": "/tmp/sample.bat", "name": "sample.bat", } self.client.sdk.get_metafields.return_value = [ _metafield("private", True), _metafield("snapshot", True), _metafield("timeout", 120), ] with patch.object(integration.Path, "open", mock_open(read_data=b"sample")): results = integration.threatzone_sandbox_upload_sample( self.client, {"entry_id": "1", "environment": "w10_x64", "private": "true"}, ) assert results[0].outputs["UUID"] == "sb-uuid" self.client.sdk.create_sandbox_submission.assert_called_once() call_kwargs = self.client.sdk.create_sandbox_submission.call_args.kwargs assert call_kwargs["environment"] == "w10_x64" assert call_kwargs["private"] is True assert call_kwargs["metafields"] == {"snapshot": True, "timeout": 120} class TestStaticAndCdrUpload(unittest.TestCase): def setUp(self): self.client = _make_client() self.client.sdk.get_user_info.return_value = _user_info() self.client.sdk.create_static_submission.return_value = SubmissionCreated(uuid="static-uuid", message="ok") self.client.sdk.create_cdr_submission.return_value = SubmissionCreated(uuid="cdr-uuid", message="ok") @patch.object(integration, "demisto") def test_static_upload(self, mock_demisto): mock_demisto.getFilePath.return_value = { "path": "/tmp/generated-entry-id", "name": "original-static.exe", } with patch.object(integration.Path, "open", mock_open(read_data=b"sample")): results = integration.threatzone_static_or_cdr_upload( self.client, {"entry_id": "1", "private": "false", "extension_check": "false"}, "static", ) assert results[0].outputs["UUID"] == "static-uuid" upload_file = self.client.sdk.create_static_submission.call_args.args[0] assert upload_file.name == "original-static.exe" call_kwargs = self.client.sdk.create_static_submission.call_args.kwargs assert call_kwargs["private"] is False assert call_kwargs["dynamic_mimetype_check"] is False @patch.object(integration, "demisto") def test_cdr_upload(self, mock_demisto): mock_demisto.getFilePath.return_value = { "path": "/tmp/generated-entry-id", "name": "original-document.docx", } with patch.object(integration.Path, "open", mock_open(read_data=b"sample")): results = integration.threatzone_static_or_cdr_upload( self.client, {"entry_id": "1", "private": "true", "extension_check": "true"}, "cdr", ) assert results[0].outputs["UUID"] == "cdr-uuid" upload_file = self.client.sdk.create_cdr_submission.call_args.args[0] assert upload_file.name == "original-document.docx" call_kwargs = self.client.sdk.create_cdr_submission.call_args.kwargs assert call_kwargs["private"] is True assert call_kwargs["dynamic_mimetype_check"] is True class TestSectionHandlers(unittest.TestCase): def setUp(self): self.client = _make_client() def test_get_indicator_result(self): self.client.sdk.get_indicators.return_value = IndicatorsResponse( items=[ Indicator( id="ind-1", name="Suspicious behavior", description="desc", category=["cat"], level="suspicious", score=50, pids=[], attackCodes=[], eventIds=[], syscallLineNumbers=[], author="system", ) ], total=1, levels=IndicatorLevels(malicious=0, suspicious=1, benign=0), ) result = integration.threatzone_get_indicator_result( self.client, { "uuid": "u", "level": "suspicious", "category": "cat", "pid": "7", "attack_code": "T1055", }, )[0] assert result.outputs["UUID"] == "u" assert result.outputs["Data"][0]["name"] == "Suspicious behavior" self.client.sdk.get_indicators.assert_called_once_with( "u", page=1, limit=integration.REPORT_FINDINGS_PAGE_SIZE, level="suspicious", category="cat", pid=7, attack_code="T1055", ) def test_get_ioc_result_fetches_every_page(self): self.client.sdk.get_iocs.side_effect = [ IoCsResponse( items=[IoC(type="domain", value="evil.example", artifacts=[])], total=2, ), IoCsResponse( items=[IoC(type="ip", value="192.0.2.1", artifacts=[])], total=2, ), ] result = integration.threatzone_get_ioc_result(self.client, {"uuid": "u", "type": "domain"})[0] assert result.outputs["Data"][0]["value"] == "evil.example" assert result.outputs["Data"][1]["value"] == "192.0.2.1" assert self.client.sdk.get_iocs.call_args_list == [ call("u", page=1, limit=integration.REPORT_FINDINGS_PAGE_SIZE, type="domain"), call("u", page=2, limit=integration.REPORT_FINDINGS_PAGE_SIZE, type="domain"), ] def test_get_yara_result(self): self.client.sdk.get_yara_rules.return_value = YaraRulesResponse( items=[YaraRule(rule="EvilRule", category="malicious", artifacts=[])], total=1, ) result = integration.threatzone_get_yara_result(self.client, {"uuid": "u", "category": "malicious"})[0] assert result.outputs["Data"][0]["rule"] == "EvilRule" self.client.sdk.get_yara_rules.assert_called_once_with( "u", page=1, limit=integration.REPORT_FINDINGS_PAGE_SIZE, category="malicious", ) def test_get_artifact_result(self): self.client.sdk.get_artifacts.return_value = ArtifactsResponse( items=[ Artifact( id="art-1", filename="dropped.bin", size=10, type="dropped_file", source="dropped", hashes=ArtifactHashes(md5="m", sha1="s", sha256="x"), tags=[], ) ], total=1, ) result = integration.threatzone_get_artifact_result(self.client, {"uuid": "u"})[0] assert result.outputs["Data"][0]["filename"] == "dropped.bin" def test_get_config_empty(self): self.client.sdk.get_extracted_configs.return_value = ExtractedConfigsResponse(items=[], total=0) result = integration.threatzone_get_config_result(self.client, {"uuid": "u"})[0] assert result.outputs is None def test_requires_uuid(self): for handler in ( integration.threatzone_get_indicator_result, integration.threatzone_get_ioc_result, integration.threatzone_get_yara_result, integration.threatzone_get_artifact_result, integration.threatzone_get_config_result, ): with pytest.raises(DemistoException): handler(self.client, {}) class TestConfigurationAndSubmissionCommands(unittest.TestCase): def setUp(self): self.client = _make_client() def test_configuration_commands(self): metafield = _metafield("timeout", 120) self.client.sdk.get_metafields.return_value = [metafield] self.client.sdk.get_environments.return_value = [] self.client.sdk.list_network_configs.return_value = [] metafields = integration.threatzone_get_metafields(self.client, {"scan_type": "sandbox"})[0] environments = integration.threatzone_get_environments(self.client, {})[0] network_configs = integration.threatzone_list_network_configs(self.client, {})[0] assert metafields.outputs == { "Data": [metafield.model_dump(by_alias=True, exclude_none=True, mode="json")], "ScanType": "sandbox", } assert environments.outputs == {"Data": []} assert network_configs.outputs == {"Data": []} self.client.sdk.get_metafields.assert_called_once_with("sandbox") self.client.sdk.get_environments.assert_called_once_with() self.client.sdk.list_network_configs.assert_called_once_with() def test_get_all_metafields_omits_filter(self): self.client.sdk.get_metafields.return_value = MagicMock() integration.threatzone_get_metafields(self.client, {}) self.client.sdk.get_metafields.assert_called_once_with() def test_get_metafields_rejects_unknown_scan_type(self): with pytest.raises(DemistoException, match="scan_type"): integration.threatzone_get_metafields(self.client, {"scan_type": "unknown"}) def test_open_in_browser_maps_sdk_arguments(self): self.client.sdk.get_user_info.return_value = _user_info() self.client.sdk.create_open_in_browser_submission.return_value = SubmissionCreated(uuid="browser-u", message="ok") result = integration.threatzone_open_in_browser( self.client, { "url": "https://example.com", "environment": "w11_x64", "auto": "false", "metafields": '{"timeout":120}', "private": "false", "configurations": '{"networkConfig":"config-id"}', }, )[0] assert result.outputs["UUID"] == "browser-u" self.client.sdk.create_open_in_browser_submission.assert_called_once_with( "https://example.com", environment="w11_x64", auto_select_environment=False, metafields={"timeout": 120}, private=False, configurations={"networkConfig": "config-id"}, ) def test_list_submissions_maps_filters(self): response = MagicMock() response.model_dump.return_value = { "items": [], "total": 0, "page": 2, "limit": 50, "totalPages": 0, } self.client.sdk.list_submissions.return_value = response result = integration.threatzone_list_submissions( self.client, { "page": "2", "limit": "50", "level": "malicious,suspicious", "tags": "tag-1,tag-2", "private": "true", "type": "file", }, )[0] assert result.outputs_prefix == "ThreatZone.Submission.List" assert result.outputs["page"] == 2 call_kwargs = self.client.sdk.list_submissions.call_args.kwargs assert call_kwargs["level"] == ["malicious", "suspicious"] assert call_kwargs["tags"] == ["tag-1", "tag-2"] assert call_kwargs["private"] is True assert call_kwargs["type"] == "file" def test_search_submissions_serializes_empty_result(self): self.client.sdk.search_by_sha256.return_value = [] result = integration.threatzone_search_submissions(self.client, {"sha256": "a" * 64})[0] assert result.outputs == {"Data": []} self.client.sdk.search_by_sha256.assert_called_once_with("a" * 64) @pytest.mark.parametrize( ("sdk_method", "section"), [ ("get_overview_summary", "OverviewSummary"), ("get_eml_analysis", "EMLAnalysis"), ("get_mitre_techniques", "MITRE"), ("get_static_scan_results", "StaticScan"), ("get_cdr_results", "CDRResult"), ("get_signature_check_results", "SignatureCheck"), ("get_processes", "Processes"), ("get_process_tree", "ProcessTree"), ("get_url_analysis", "URLAnalysis"), ("get_network_summary", "NetworkSummary"), ], ) def test_uuid_section_sdk_mappings(sdk_method, section): client = _make_client() getattr(client.sdk, sdk_method).return_value = {"value": sdk_method} result = integration.threatzone_get_uuid_section(client, {"uuid": "u"}, sdk_method, section, "Title")[0] assert result.outputs == {"UUID": "u", "Data": {"value": sdk_method}} getattr(client.sdk, sdk_method).assert_called_once_with("u") class TestTelemetryCommands(unittest.TestCase): def setUp(self): self.client = _make_client() def test_behaviours_defaults_and_filters(self): self.client.sdk.get_behaviours.return_value = {"items": [], "total": 0} result = integration.threatzone_get_behaviours( self.client, {"uuid": "u", "pid": "42", "process_name": "sample.exe"}, )[0] assert result.outputs["UUID"] == "u" self.client.sdk.get_behaviours.assert_called_once_with( "u", type=None, pid=42, operation=None, process_name="sample.exe", page=1, limit=100, ) def test_behaviours_bounds_limit(self): with pytest.raises(DemistoException, match="between 1 and 500"): integration.threatzone_get_behaviours(self.client, {"uuid": "u", "limit": "501"}) def test_syscalls_defaults_and_limit_bound(self): self.client.sdk.get_syscalls.return_value = {"items": [], "total": 0} integration.threatzone_get_syscalls(self.client, {"uuid": "u"}) self.client.sdk.get_syscalls.assert_called_once_with("u", page=1, limit=500) with pytest.raises(DemistoException, match="between 1 and 2000"): integration.threatzone_get_syscalls(self.client, {"uuid": "u", "limit": "2001"}) def test_network_window_mappings(self): for sdk_method in ( "get_dns_queries", "get_http_requests", "get_tcp_connections", "get_udp_connections", "get_network_threats", ): sdk_mock = getattr(self.client.sdk, sdk_method) sdk_mock.return_value = [] result = integration.threatzone_get_network_data( self.client, {"uuid": "u", "limit": "1000", "skip": "0"}, sdk_method, "Section", "Title", )[0] assert result.outputs == {"UUID": "u", "Data": []} sdk_mock.assert_called_once_with("u", limit=1000, skip=0) def test_network_window_rejects_out_of_range(self): with pytest.raises(DemistoException, match="between 0 and 1000"): integration.threatzone_get_network_data( self.client, {"uuid": "u", "skip": "1001"}, "get_dns_queries", "DNSQueries", "DNS Queries", ) class TestGetResult(unittest.TestCase): def setUp(self): self.client = _make_client() self.client.sdk.get_submission.return_value = _submission() # By default, the legacy IOC query returns no items. self.client.sdk.get_iocs.return_value = IoCsResponse(items=[], total=0) def test_basic_result(self): results = integration.threatzone_get_result(self.client, {"uuid": "u"}) prefixes = [r.outputs_prefix for r in results if hasattr(r, "outputs_prefix")] assert "ThreatZone.Submission" in prefixes assert "ThreatZone.Analysis" in prefixes assert "ThreatZone.IOC" in prefixes analysis = next(r for r in results if r.outputs_prefix == "ThreatZone.Analysis") submission = next(r for r in results if r.outputs_prefix == "ThreatZone.Submission") assert analysis.outputs["LEVEL"] == 3 assert analysis.outputs["STATUS"] == 5 assert analysis.outputs["REPORT"]["status"] == 5 assert submission.outputs["Summary"]["REPORT"]["status"] == 5 assert submission.outputs["reports"][0]["status"] == "completed" assert analysis.outputs["SHA256"].startswith("6e899ff7") def test_legacy_iocs_include_every_page(self): self.client.sdk.get_iocs.side_effect = [ IoCsResponse(items=[IoC(type="domain", value="one.example", artifacts=[])], total=2), IoCsResponse(items=[IoC(type="domain", value="two.example", artifacts=[])], total=2), ] results = integration.threatzone_get_result(self.client, {"uuid": "u"}) legacy_iocs = next(result for result in results if result.outputs_prefix == "ThreatZone.IOC") assert legacy_iocs.outputs["DOMAIN"] == ["one.example", "two.example"] def test_details_use_paginated_finding_endpoints(self): self.client.sdk.get_indicators.return_value = IndicatorsResponse( items=[], total=0, levels=IndicatorLevels(malicious=0, suspicious=0, benign=0), ) self.client.sdk.get_yara_rules.return_value = YaraRulesResponse(items=[], total=0) self.client.sdk.get_artifacts.return_value = ArtifactsResponse(items=[], total=0) self.client.sdk.get_extracted_configs.return_value = ExtractedConfigsResponse(items=[], total=0) integration.threatzone_get_result(self.client, {"uuid": "u", "details": "true"}) expected_page_call = call("c89d310b-7862-4534-998a-3eb39d9a9d42", page=1, limit=100) assert self.client.sdk.get_indicators.call_args == expected_page_call assert self.client.sdk.get_yara_rules.call_args == expected_page_call assert self.client.sdk.get_iocs.call_args_list == [expected_page_call] def test_url_analysis_preserves_legacy_type_label(self): self.client.sdk.get_submission.return_value = _submission(report_type="url_analysis") results = integration.threatzone_get_result(self.client, {"uuid": "u"}) analysis = next(result for result in results if result.outputs_prefix == "ThreatZone.Analysis") assert analysis.outputs["TYPE"] == "urlAnalysis" def test_declined_status_raises(self): self.client.sdk.get_submission.return_value = _submission(status="error") with pytest.raises(DemistoException) as excinfo: integration.threatzone_get_result(self.client, {"uuid": "u"}) assert "declined" in str(excinfo.value).lower() class TestDownloads(unittest.TestCase): def setUp(self): self.client = _make_client() @patch.object(integration, "_save_download", return_value={"EntryID": "entry-1"}) def test_download_html_report(self, save_download_mock): download = MagicMock() self.client.sdk.download_html_report.return_value = download result = integration.threatzone_get_html_report_file(self.client, {"uuid": "u"}) assert result["EntryID"] == "entry-1" save_download_mock.assert_called_once_with(download, "threatzone-report-u.html") @patch.object(integration, "_save_download", return_value={"EntryID": "entry-2"}) def test_download_cdr_result(self, save_download_mock): download = MagicMock() self.client.sdk.download_cdr_result.return_value = download result = integration.threatzone_get_sanitized_file(self.client, {"uuid": "u"}) assert result["EntryID"] == "entry-2" save_download_mock.assert_called_once_with(download, "sanitized-u.zip") @patch.object(integration, "file_result_existing_file", return_value={"EntryID": "entry-3"}) def test_save_download_streams_to_existing_file(self, existing_file_result_mock): download = MagicMock() download.filename = "report.html" download.save.return_value = Path("report.html") result = integration._save_download(download, "fallback.html") assert result["EntryID"] == "entry-3" download.save.assert_called_once_with("report.html") download.read.assert_not_called() download.close.assert_called_once_with() existing_file_result_mock.assert_called_once_with("report.html", "report.html") @patch.object(integration, "file_result_existing_file", return_value={"EntryID": "entry-4"}) def test_save_download_uses_fallback_for_unnamed_response(self, existing_file_result_mock): download = MagicMock() download.filename = "download" download.save.return_value = Path("fallback.html") result = integration._save_download(download, "fallback.html") assert result["EntryID"] == "entry-4" download.save.assert_called_once_with("fallback.html") download.close.assert_called_once_with() existing_file_result_mock.assert_called_once_with("fallback.html", "fallback.html") @patch.object(integration, "_save_download", return_value={"EntryID": "entry"}) def test_sdk_stream_download_mappings(self, save_download_mock): mappings = ( ("get_static_scan_strings", "{uuid}_strings.json", None), ("download_sample", "sample-{uuid}", None), ("download_artifact", "artifact-{uuid}", "artifact_id"), ("download_pcap", "threatzone-{uuid}.pcap", None), ) for sdk_method, fallback, id_argument in mappings: save_download_mock.reset_mock() download = MagicMock() getattr(self.client.sdk, sdk_method).return_value = download args = {"uuid": "u", "artifact_id": "artifact-1"} integration.threatzone_download_sdk_file( self.client, args, sdk_method, fallback, id_argument=id_argument, ) expected_args = ("u", "artifact-1") if id_argument else ("u",) getattr(self.client.sdk, sdk_method).assert_called_once_with(*expected_args) save_download_mock.assert_called_once_with(download, fallback.format(uuid="u")) @patch.object(integration, "_save_download", return_value={"EntryID": "yara-entry"}) def test_generated_yara_immediate_success(self, save_download_mock): download = MagicMock() self.client.sdk.download_yara_rule.return_value = download result = integration.threatzone_download_yara_rule(self.client, {"uuid": "u"}) assert result == {"EntryID": "yara-entry"} save_download_mock.assert_called_once_with(download, "u.yar") @patch.object(integration, "_save_download", return_value={"EntryID": "yara-entry"}) @patch.object(integration.demisto, "executeCommand") @patch.object(integration.time, "monotonic", side_effect=[0.0, 1.0, 4.0]) def test_generated_yara_polls_with_server_retry(self, monotonic_mock, execute_command_mock, save_download_mock): download = MagicMock() self.client.sdk.download_yara_rule.side_effect = [ YaraRulePendingError("pending", retry_after=2.5), YaraRulePendingError("pending"), download, ] integration.threatzone_download_yara_rule(self.client, {"uuid": "u", "timeout": "120"}) assert execute_command_mock.call_args_list == [ call("Sleep", {"seconds": "2.5"}), call("Sleep", {"seconds": str(integration.YARA_POLL_INTERVAL_SECONDS)}), ] save_download_mock.assert_called_once_with(download, "u.yar") @patch.object(integration.demisto, "executeCommand") @patch.object(integration.time, "monotonic", side_effect=[0.0, 119.0]) def test_generated_yara_timeout(self, monotonic_mock, execute_command_mock): self.client.sdk.download_yara_rule.side_effect = YaraRulePendingError("pending", retry_after=2.0) with pytest.raises(DemistoException, match="Timed out after 120 seconds"): integration.threatzone_download_yara_rule(self.client, {"uuid": "u"}) execute_command_mock.assert_not_called() @patch.object(integration.demisto, "executeCommand") @patch.object(integration.time, "monotonic", side_effect=[0.0, 0.0, 0.0]) def test_generated_yara_polling_has_an_attempt_limit(self, monotonic_mock, execute_command_mock): self.client.sdk.download_yara_rule.side_effect = YaraRulePendingError("pending", retry_after=0.0) with pytest.raises(DemistoException, match="Timed out after 2 seconds"): integration.threatzone_download_yara_rule(self.client, {"uuid": "u", "timeout": "2"}) assert self.client.sdk.download_yara_rule.call_count == 2 assert execute_command_mock.call_args_list == [ call("Sleep", {"seconds": str(integration.MIN_YARA_POLL_INTERVAL_SECONDS)}), call("Sleep", {"seconds": str(integration.MIN_YARA_POLL_INTERVAL_SECONDS)}), ] @patch.object(integration, "fileResult", return_value={"EntryID": "screenshot"}) def test_screenshot_bytes_become_war_room_file(self, file_result_mock): self.client.sdk.get_screenshot.return_value = b"png" result = integration.threatzone_download_url_screenshot(self.client, {"uuid": "u"}) assert result == {"EntryID": "screenshot"} file_result_mock.assert_called_once_with("threatzone-url-screenshot-u.png", b"png") @patch.object(integration, "fileResult", return_value={"EntryID": "media"}) def test_media_file_uses_validated_server_filename(self, file_result_mock): media = MagicMock() media.id = "media-1" media.name = "screen.png" self.client.sdk.list_media_files.return_value = [media] self.client.sdk.get_media_file.return_value = b"png" result = integration.threatzone_download_media_file( self.client, {"uuid": "u", "file_id": "media-1"}, ) assert result == {"EntryID": "media"} file_result_mock.assert_called_once_with("screen.png", b"png") def test_media_file_rejects_unsafe_server_filename(self): media = MagicMock() media.id = "media-1" media.name = "../screen.png" self.client.sdk.list_media_files.return_value = [media] with pytest.raises(DemistoException, match="unsafe media filename"): integration.threatzone_download_media_file( self.client, {"uuid": "u", "file_id": "media-1"}, ) class TestSdkExceptionFormatting(unittest.TestCase): def test_known_exceptions(self): assert "Authorization" in integration._format_sdk_exception(AuthenticationError("nope", status_code=401)) assert "not found" in integration._format_sdk_exception(NotFoundError("nope", status_code=404)).lower() assert "API error" in integration._format_sdk_exception(APIError("boom", status_code=500)) assert "timed out" in integration._format_sdk_exception(AnalysisTimeoutError("late", uuid="u", elapsed=1.0)) assert "not yet available" in integration._format_sdk_exception(ReportUnavailableError("wait")) class TestTestModule(unittest.TestCase): def test_ok(self): client = _make_client() client.sdk.get_user_info.return_value = _user_info() assert integration.test_module(client) == "ok" def test_auth_failure(self): client = _make_client() client.sdk.get_user_info.side_effect = AuthenticationError("bad", status_code=401) result = integration.test_module(client) assert "Authorization" in result class TestSdkConsumerContract(unittest.TestCase): def test_completed_submission_through_fake_api(self): sample_bytes = b"threatzone-sdk-contract-sample" sample_sha256 = hashlib.sha256(sample_bytes).hexdigest() fake_api = FakeThreatZoneAPI() scenarios.seed_malicious_pe(fake_api, sha256=sample_sha256) sdk = ThreatZoneSDK( api_key="test-key", base_url="https://fake.threat.zone/public-api", http_client=fake_api.as_httpx_client(), ) created = sdk.create_sandbox_submission(sample_bytes, private=True) sdk.get_submission(created.uuid) sdk.get_submission(created.uuid) client = _make_client() client.sdk = sdk results = integration.threatzone_get_result(client, {"uuid": created.uuid}) analysis = next(result for result in results if result.outputs_prefix == "ThreatZone.Analysis") assert analysis.outputs["STATUS"] == 5 assert analysis.outputs["LEVEL"] == 3 assert analysis.outputs["REPORT"]["status"] == 5 @patch.object(integration, "_save_download", return_value={"EntryID": "fake-download"}) def test_extended_commands_through_fake_api(self, save_download_mock): fake_api = FakeThreatZoneAPI() scenarios.seed_malicious_pe(fake_api) sdk = ThreatZoneSDK( api_key="test-key", base_url="https://fake.threat.zone/public-api", http_client=fake_api.as_httpx_client(), ) created = sdk.create_sandbox_submission(b"sample", private=True) sdk.get_submission(created.uuid) sdk.get_submission(created.uuid) client = _make_client() client.sdk = sdk assert integration.threatzone_get_metafields(client, {"scan_type": "sandbox"})[0].outputs["Data"] assert integration.threatzone_list_submissions(client, {})[0].outputs["items"] assert integration.threatzone_get_uuid_section( client, {"uuid": created.uuid}, "get_overview_summary", "OverviewSummary", "Overview", )[0].outputs["Data"] assert integration.threatzone_get_network_data( client, {"uuid": created.uuid}, "get_network_threats", "NetworkThreats", "Network Threats", )[0].outputs["Data"] integration.threatzone_download_sdk_file( client, {"uuid": created.uuid}, "download_sample", "sample-{uuid}", ) integration.threatzone_download_yara_rule(client, {"uuid": created.uuid}) assert save_download_mock.call_count == 2 for download_call in save_download_mock.call_args_list: download_call.args[0].close() if __name__ == "__main__": unittest.main()