diff --git a/tests/api.py b/tests/api.py index 3bc8e4b..84038cd 100644 --- a/tests/api.py +++ b/tests/api.py @@ -13,7 +13,7 @@ def setUp(self): @patch("yeti.api.requests.Session.post") def test_auth_api_key(self, mock_post): mock_response = MagicMock() - mock_response.bytes = b'{"access_token": "fake_token"}' + mock_response.content = b'{"access_token": "fake_token"}' mock_post.return_value = mock_response self.api.auth_api_key("fake_apikey") @@ -26,7 +26,7 @@ def test_auth_api_key(self, mock_post): @patch("yeti.api.requests.Session.post") def test_search_indicators(self, mock_post): mock_response = MagicMock() - mock_response.bytes = b'{"indicators": [{"name": "test"}]}' + mock_response.content = b'{"indicators": [{"name": "test"}]}' mock_post.return_value = mock_response result = self.api.search_indicators(name="test") @@ -39,7 +39,7 @@ def test_search_indicators(self, mock_post): @patch("yeti.api.requests.Session.post") def test_search_entities(self, mock_post): mock_response = MagicMock() - mock_response.bytes = b'{"entities": [{"name": "test_entity"}]}' + mock_response.content = b'{"entities": [{"name": "test_entity"}]}' mock_post.return_value = mock_response result = self.api.search_entities(name="test_entity") @@ -52,7 +52,7 @@ def test_search_entities(self, mock_post): @patch("yeti.api.requests.Session.post") def test_search_observables(self, mock_post): mock_response = MagicMock() - mock_response.bytes = b'{"observables": [{"value": "test_value"}]}' + mock_response.content = b'{"observables": [{"value": "test_value"}]}' mock_post.return_value = mock_response result = self.api.search_observables(value="test_value") @@ -65,8 +65,8 @@ def test_search_observables(self, mock_post): @patch("yeti.api.requests.Session.post") def test_new_entity(self, mock_post): mock_response = MagicMock() - mock_response.bytes = b'{"id": "new_entity"}' - mock_response.bytes = b'{"id": "new_entity"}' + mock_response.content = b'{"id": "new_entity"}' + mock_response.content = b'{"id": "new_entity"}' mock_post.return_value = mock_response result = self.api.new_entity({"name": "test_entity"}) @@ -79,7 +79,7 @@ def test_new_entity(self, mock_post): @patch("yeti.api.requests.Session.post") def test_new_indicator(self, mock_post): mock_response = MagicMock() - mock_response.bytes = b'{"id": "new_indicator"}' + mock_response.content = b'{"id": "new_indicator"}' mock_post.return_value = mock_response result = self.api.new_indicator({"name": "test_indicator"}) @@ -92,7 +92,7 @@ def test_new_indicator(self, mock_post): @patch("yeti.api.requests.Session.patch") def test_patch_indicator(self, mock_patch): mock_response = MagicMock() - mock_response.bytes = b'{"id": "patched_indicator"}' + mock_response.content = b'{"id": "patched_indicator"}' mock_patch.return_value = mock_response result = self.api.patch_indicator(1, {"name": "patched_indicator"}) @@ -105,7 +105,7 @@ def test_patch_indicator(self, mock_patch): @patch("yeti.api.requests.Session.post") def test_search_dfiq(self, mock_post): mock_response = MagicMock() - mock_response.bytes = b'{"dfiq": [{"name": "test_dfiq"}]}' + mock_response.content = b'{"dfiq": [{"name": "test_dfiq"}]}' mock_post.return_value = mock_response result = self.api.search_dfiq(name="test_dfiq") @@ -118,7 +118,7 @@ def test_search_dfiq(self, mock_post): @patch("yeti.api.requests.Session.post") def test_new_dfiq_from_yaml(self, mock_post): mock_response = MagicMock() - mock_response.bytes = b'{"id": "new_dfiq"}' + mock_response.content = b'{"id": "new_dfiq"}' mock_post.return_value = mock_response result = self.api.new_dfiq_from_yaml("type", "yaml_content") @@ -135,7 +135,7 @@ def test_new_dfiq_from_yaml(self, mock_post): @patch("yeti.api.requests.Session.patch") def test_patch_dfiq_from_yaml(self, mock_patch): mock_response = MagicMock() - mock_response.bytes = b'{"id": "patched_dfiq"}' + mock_response.content = b'{"id": "patched_dfiq"}' mock_patch.return_value = mock_response result = self.api.patch_dfiq_from_yaml("type", "yaml_content", 1) @@ -152,7 +152,7 @@ def test_patch_dfiq_from_yaml(self, mock_patch): @patch("yeti.api.requests.Session.post") def test_download_dfiq_archive(self, mock_post): mock_response = MagicMock() - mock_response.bytes = b"archive_content" + mock_response.content = b"archive_content" mock_post.return_value = mock_response result = self.api.download_dfiq_archive() @@ -165,7 +165,7 @@ def test_download_dfiq_archive(self, mock_post): @patch("yeti.api.requests.Session.post") def test_upload_dfiq_archive(self, mock_post): mock_response = MagicMock() - mock_response.bytes = b'{"uploaded": 1}' + mock_response.content = b'{"uploaded": 1}' mock_post.return_value = mock_response with patch("builtins.open", unittest.mock.mock_open(read_data=b"data")): @@ -182,7 +182,7 @@ def test_upload_dfiq_archive(self, mock_post): @patch("yeti.api.requests.Session.post") def test_add_observable(self, mock_post): mock_response = MagicMock() - mock_response.bytes = b'{"id": "new_observable"}' + mock_response.content = b'{"id": "new_observable"}' mock_post.return_value = mock_response result = self.api.add_observable("value", "type") @@ -195,7 +195,7 @@ def test_add_observable(self, mock_post): @patch("yeti.api.requests.Session.post") def test_add_observables_bulk(self, mock_post): mock_response = MagicMock() - mock_response.bytes = b'{"added": [], "failed": []}' + mock_response.content = b'{"added": [], "failed": []}' mock_post.return_value = mock_response result = self.api.add_observables_bulk([{"value": "value", "type": "type"}]) @@ -208,7 +208,7 @@ def test_add_observables_bulk(self, mock_post): @patch("yeti.api.requests.Session.post") def test_tag_object(self, mock_post): mock_response = MagicMock() - mock_response.bytes = b'{"id": "tagged_object"}' + mock_response.content = b'{"id": "tagged_object"}' mock_post.return_value = mock_response result = self.api.tag_object({"id": "1", "root_type": "indicator"}, ["tag1"]) @@ -221,7 +221,7 @@ def test_tag_object(self, mock_post): @patch("yeti.api.requests.Session.post") def test_link_objects(self, mock_post): mock_response = MagicMock() - mock_response.bytes = b'{"id": "link"}' + mock_response.content = b'{"id": "link"}' mock_post.return_value = mock_response result = self.api.link_objects( @@ -243,7 +243,7 @@ def test_link_objects(self, mock_post): @patch("yeti.api.requests.Session.post") def test_search_graph(self, mock_post): mock_response = MagicMock() - mock_response.bytes = b'{"graph": "data"}' + mock_response.content = b'{"graph": "data"}' mock_post.return_value = mock_response result = self.api.search_graph("source", "graph", ["type"]) diff --git a/yeti/api.py b/yeti/api.py index 0cc64b4..1533a91 100644 --- a/yeti/api.py +++ b/yeti/api.py @@ -3,10 +3,11 @@ import json from typing import Any, Sequence -import yeti.errors as errors import requests import requests_toolbelt.multipart.encoder as encoder +import yeti.errors as errors + TYPE_TO_ENDPOINT = { "indicator": "/api/v2/indicators", "entity": "/api/v2/entities", @@ -86,7 +87,7 @@ def do_request( except requests.exceptions.HTTPError as e: raise errors.YetiApiError(e.response.status_code, e.response.text) - return response.bytes + return response.content def auth_api_key(self, apikey: str) -> None: """Authenticates a session using an API key."""