From 9d50278ae0029f91731b688da5cb3de6230fbddf Mon Sep 17 00:00:00 2001 From: Thomas Chopitea Date: Thu, 1 May 2025 21:00:37 +0000 Subject: [PATCH 1/8] Add functions to find() objects by name / type --- tests/api.py | 88 ++++++++++++++++++++++++++++++++++++++++++++++++ yeti/api.py | 95 +++++++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 182 insertions(+), 1 deletion(-) diff --git a/tests/api.py b/tests/api.py index 4eba3d9..2fa8332 100644 --- a/tests/api.py +++ b/tests/api.py @@ -332,6 +332,94 @@ def test_get_yara_bundle_with_overlays(self, mock_post): }, ) + @patch("yeti.api.requests.Session.get") + def test_find_indicator(self, mock_get): + mock_response = MagicMock() + mock_response.content = b'{"id": "found_indicator"}' + mock_get.return_value = mock_response + + result = self.api.find_indicator(name="test_indicator", type="indicator.test") + self.assertEqual(result, {"id": "found_indicator"}) + mock_get.assert_called_with( + "http://fake-url/api/v2/indicators/?name=test_indicator&type=indicator.test", + ) + + # Test 404 case + mock_exception_with_status_code = requests.exceptions.HTTPError() + mock_exception_with_status_code.response = MagicMock() + mock_exception_with_status_code.response.status_code = 404 + mock_response.raise_for_status.side_effect = mock_exception_with_status_code + mock_get.return_value = mock_response + + result = self.api.find_indicator(name="not_found", type="indicator.test") + self.assertIsNone(result) + + @patch("yeti.api.requests.Session.get") + def test_find_entity(self, mock_get): + mock_response = MagicMock() + mock_response.content = b'{"id": "found_entity"}' + mock_get.return_value = mock_response + + result = self.api.find_entity(name="test_entity", type="entity.test") + self.assertEqual(result, {"id": "found_entity"}) + mock_get.assert_called_with( + "http://fake-url/api/v2/entities/?name=test_entity&type=entity.test", + ) + + # Test 404 case + mock_exception_with_status_code = requests.exceptions.HTTPError() + mock_exception_with_status_code.response = MagicMock() + mock_exception_with_status_code.response.status_code = 404 + mock_response.raise_for_status.side_effect = mock_exception_with_status_code + mock_get.return_value = mock_response + + result = self.api.find_entity(name="not_found", type="entity.test") + self.assertIsNone(result) + + @patch("yeti.api.requests.Session.get") + def test_find_observable(self, mock_get): + mock_response = MagicMock() + mock_response.content = b'{"id": "found_observable"}' + mock_get.return_value = mock_response + + result = self.api.find_observable(value="test_value", type="observable.test") + self.assertEqual(result, {"id": "found_observable"}) + mock_get.assert_called_with( + "http://fake-url/api/v2/observables/?value=test_value&type=observable.test", + ) + + # Test 404 case + mock_exception_with_status_code = requests.exceptions.HTTPError() + mock_exception_with_status_code.response = MagicMock() + mock_exception_with_status_code.response.status_code = 404 + mock_response.raise_for_status.side_effect = mock_exception_with_status_code + mock_get.return_value = mock_response + + result = self.api.find_observable(value="not_found", type="observable.test") + self.assertIsNone(result) + + @patch("yeti.api.requests.Session.get") + def test_find_dfiq(self, mock_get): + mock_response = MagicMock() + mock_response.content = b'{"id": "found_dfiq"}' + mock_get.return_value = mock_response + + result = self.api.find_dfiq(name="test_dfiq", dfiq_type="scenario") + self.assertEqual(result, {"id": "found_dfiq"}) + mock_get.assert_called_with( + "http://fake-url/api/v2/dfiq/?name=test_dfiq&type=scenario", + ) + + # Test 404 case + mock_exception_with_status_code = requests.exceptions.HTTPError() + mock_exception_with_status_code.response = MagicMock() + mock_exception_with_status_code.response.status_code = 404 + mock_response.raise_for_status.side_effect = mock_exception_with_status_code + mock_get.return_value = mock_response + + result = self.api.find_dfiq(name="not_found", dfiq_type="scenario") + self.assertIsNone(result) + if __name__ == "__main__": unittest.main() diff --git a/yeti/api.py b/yeti/api.py index c41a64b..e133ca3 100644 --- a/yeti/api.py +++ b/yeti/api.py @@ -2,6 +2,7 @@ import json import logging +import urllib.parse from typing import Any, Sequence import requests @@ -63,6 +64,7 @@ def do_request( body: bytes | None = None, headers: dict[str, Any] | None = None, retries: int = 3, + params: dict[str, Any] | None = None, ) -> bytes: """Issues a request to the given URL. @@ -73,6 +75,7 @@ def do_request( body: The body to include in the request. headers: Extra headers to include in the request. retries: The number of times to retry the request. + params: The query parameters to include in the request. Returns: The response from the API; a bytes object. @@ -90,6 +93,8 @@ def do_request( request_kwargs["json"] = json_data if body: request_kwargs["body"] = body + if params: + url = f"{url}?{urllib.parse.urlencode(params)}" try: if method == "POST": @@ -145,6 +150,28 @@ def refresh_auth(self): else: logger.warning("No auth function set, cannot refresh auth.") + def find_indicator(self, name: str, type: str) -> YetiObject | None: + """Finds an indicator in Yeti by name and type. + + Args: + name: The name of the indicator to find. + type: The type of the indicator to find. + + Returns: + The response from the API; a dict representing the indicator. + """ + try: + response = self.do_request( + "GET", + f"{self._url_root}/api/v2/indicators/", + params={"name": name, "type": type}, + ) + except errors.YetiApiError as e: + if e.status_code == 404: + return None + raise + return json.loads(response) + def search_indicators( self, name: str | None = None, @@ -163,7 +190,7 @@ def search_indicators( tags: The tags of the indicator to search for. Returns: - The response from the API; a dict representing the indicator. + The response from the API; a list of dicts representing indicators. """ if not any([name, indicator_type, pattern, tags]): @@ -188,6 +215,28 @@ def search_indicators( ) return json.loads(response)["indicators"] + def find_entity(self, name: str, type: str) -> YetiObject | None: + """Finds an entity in Yeti by name. + + Args: + name: The name of the entity to find. + type: The type of the entity to find. + + Returns: + The response from the API; a dict representing the entity. + """ + try: + response = self.do_request( + "GET", + f"{self._url_root}/api/v2/entities/", + params={"name": name, "type": type}, + ) + except errors.YetiApiError as e: + if e.status_code == 404: + return None + raise + return json.loads(response) + def search_entities(self, name: str) -> list[YetiObject]: params = {"query": {"name": name}, "count": 0} response = self.do_request( @@ -197,6 +246,28 @@ def search_entities(self, name: str) -> list[YetiObject]: ) return json.loads(response)["entities"] + def find_observable(self, value: str, type: str) -> YetiObject | None: + """Finds an observable in Yeti by value and type. + + Args: + value: The value of the observable to find. + type: The type of the observable to find. + + Returns: + The response from the API; a dict representing the observable. + """ + try: + response = self.do_request( + "GET", + f"{self._url_root}/api/v2/observables/", + params={"value": value, "type": type}, + ) + except errors.YetiApiError as e: + if e.status_code == 404: + return None + raise + return json.loads(response) + def search_observables(self, value: str) -> list[YetiObject]: """Searches for an observable in Yeti. @@ -330,6 +401,28 @@ def get_yara_bundle_with_overlays( return json.loads(result) + def find_dfiq(self, name: str, dfiq_type: str) -> YetiObject | None: + """Finds a DFIQ in Yeti by name and type. + + Args: + name: The name of the DFIQ to find. + dfiq_type: The type of the DFIQ to find. + + Returns: + The response from the API; a dict representing the DFIQ object. + """ + try: + response = self.do_request( + "GET", + f"{self._url_root}/api/v2/dfiq/", + params={"name": name, "type": dfiq_type}, + ) + except errors.YetiApiError as e: + if e.status_code == 404: + return None + raise + return json.loads(response) + def search_dfiq(self, name: str, dfiq_type: str | None = None) -> list[YetiObject]: """Searches for a DFIQ in Yeti. From efab3948ed046f84f745dfd587b8033545486da7 Mon Sep 17 00:00:00 2001 From: Thomas Chopitea Date: Thu, 1 May 2025 21:07:51 +0000 Subject: [PATCH 2/8] Restore e2e tests --- .github/workflows/e2e.yml | 43 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) create mode 100644 .github/workflows/e2e.yml diff --git a/.github/workflows/e2e.yml b/.github/workflows/e2e.yml new file mode 100644 index 0000000..db999d8 --- /dev/null +++ b/.github/workflows/e2e.yml @@ -0,0 +1,43 @@ +name: e2e tests + +on: [pull_request] + +jobs: + unittest: + runs-on: ubuntu-latest + env: + YETI_ENDPOINT: "http://localhost:80" + strategy: + matrix: + os: [ubuntu-latest] + python-version: ["3.10"] + steps: + - uses: actions/checkout@v4 + + - name: Install dependencies + run: | + sudo apt-get update + sudo apt-get install -y python3-pip git + sudo pip3 install poetry + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + cache: poetry + + - name: Install a Yeti prod deployment + run: | + git clone https://github.com/yeti-platform/yeti-docker + cd yeti-docker/prod + ./init.sh + sleep 10 + - name: Create test Yeti user + run: | + cd yeti-docker/prod + echo "YETI_API_KEY=$(docker compose run --rm api create-user test test --admin | awk -F'test:' '{print $2}')" >> $GITHUB_ENV + - name: Install Python dependencies + run: poetry install --no-root + + - name: e2e testing + run: | + YETI_API_KEY=$YETI_API_KEY poetry run python -m unittest tests/e2e.py From 43a0015e098b5f3920571ad879737325c7b6d020 Mon Sep 17 00:00:00 2001 From: Thomas Chopitea Date: Thu, 1 May 2025 21:16:49 +0000 Subject: [PATCH 3/8] Create directory for bloomfilter in e2e tests --- .github/workflows/e2e.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/e2e.yml b/.github/workflows/e2e.yml index db999d8..e4df528 100644 --- a/.github/workflows/e2e.yml +++ b/.github/workflows/e2e.yml @@ -29,6 +29,7 @@ jobs: run: | git clone https://github.com/yeti-platform/yeti-docker cd yeti-docker/prod + mkdir -p /opt/yeti/bloomfilters ./init.sh sleep 10 - name: Create test Yeti user From 058e5c37d19244deaae7b32d02e724c24d3e6b5b Mon Sep 17 00:00:00 2001 From: Thomas Chopitea Date: Thu, 1 May 2025 22:46:19 +0000 Subject: [PATCH 4/8] Add e2e.py file --- .github/workflows/e2e.yml | 3 +- tests/e2e.py | 65 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 2 deletions(-) create mode 100644 tests/e2e.py diff --git a/.github/workflows/e2e.yml b/.github/workflows/e2e.yml index e4df528..e2fef04 100644 --- a/.github/workflows/e2e.yml +++ b/.github/workflows/e2e.yml @@ -3,7 +3,7 @@ name: e2e tests on: [pull_request] jobs: - unittest: + e2e: runs-on: ubuntu-latest env: YETI_ENDPOINT: "http://localhost:80" @@ -29,7 +29,6 @@ jobs: run: | git clone https://github.com/yeti-platform/yeti-docker cd yeti-docker/prod - mkdir -p /opt/yeti/bloomfilters ./init.sh sleep 10 - name: Create test Yeti user diff --git a/tests/e2e.py b/tests/e2e.py new file mode 100644 index 0000000..c24a5c2 --- /dev/null +++ b/tests/e2e.py @@ -0,0 +1,65 @@ +import os +import time +import unittest +from unittest.mock import MagicMock, patch + +import requests + +from yeti import errors +from yeti.api import YetiApi + + +class YetiEndToEndTest(unittest.TestCase): + def setUp(self): + self.api = YetiApi(os.getenv("YETI_ENDPOINT")) + + def test_auth_api_key(self): + self.api.auth_api_key(os.getenv("YETI_API_KEY")) + self.api.search_indicators(name="test") + + def test_no_auth(self): + with self.assertRaises(errors.YetiAuthError) as error: + self.api.search_indicators(name="test") + self.assertIn( + "401 Client Error: Unauthorized for url: ", + str(error.exception), + ) + + def test_new_indicator(self): + self.api.auth_api_key(os.getenv("YETI_API_KEY")) + indicator = self.api.new_indicator( + { + "name": "test", + "type": "regex", + "description": "test", + "pattern": "test[0-9]", + "diamond": "victim", + } + ) + self.assertEqual(indicator["name"], "test") + self.assertRegex(indicator["id"], r"[0-9]+") + + def test_auth_refresh(self): + self.api.auth_api_key(os.getenv("YETI_API_KEY")) + self.api.search_indicators(name="test") + + time.sleep(3) + + self.api.search_indicators(name="test") + + def test_search_indicators(self): + self.api.auth_api_key(os.getenv("YETI_API_KEY")) + self.api.auth_api_key(os.getenv("YETI_API_KEY")) + indicator = self.api.new_indicator( + { + "name": "testSearch", + "type": "regex", + "description": "test", + "pattern": "test[0-9]", + "diamond": "victim", + } + ) + time.sleep(5) + result = self.api.search_indicators(name="testSear") + self.assertEqual(len(result), 1, result) + self.assertEqual(result[0]["name"], "testSearch") From 00888198d4f220055c8591fb488a76cff0343816 Mon Sep 17 00:00:00 2001 From: Thomas Chopitea Date: Thu, 1 May 2025 22:46:30 +0000 Subject: [PATCH 5/8] Don't discover tests --- .github/workflows/unittests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/unittests.yml b/.github/workflows/unittests.yml index cf66077..f4cae35 100644 --- a/.github/workflows/unittests.yml +++ b/.github/workflows/unittests.yml @@ -23,4 +23,4 @@ jobs: run: poetry install --no-root - name: Test with unittest run: | - poetry run python -m unittest discover -s tests/ -p '*.py' + poetry run python -m unittest tests/api.py From 7ec7304f34ed0ffc92dadbcb2cde7e54bffaf736 Mon Sep 17 00:00:00 2001 From: Thomas Chopitea Date: Fri, 2 May 2025 09:26:03 +0000 Subject: [PATCH 6/8] Fix docstring --- yeti/api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/yeti/api.py b/yeti/api.py index e133ca3..ba375ed 100644 --- a/yeti/api.py +++ b/yeti/api.py @@ -71,7 +71,7 @@ def do_request( Args: method: The HTTP method to use. url: The URL to issue the request to. - json: The JSON payload to include in the request. + json_data: The JSON payload to include in the request. body: The body to include in the request. headers: Extra headers to include in the request. retries: The number of times to retry the request. From c64f2a7bff4b350467a46dbae857de0ae7a15835 Mon Sep 17 00:00:00 2001 From: Thomas Chopitea Date: Fri, 2 May 2025 12:25:22 +0000 Subject: [PATCH 7/8] Tweak e2e test --- tests/e2e.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/e2e.py b/tests/e2e.py index c24a5c2..6516ca0 100644 --- a/tests/e2e.py +++ b/tests/e2e.py @@ -50,7 +50,7 @@ def test_auth_refresh(self): def test_search_indicators(self): self.api.auth_api_key(os.getenv("YETI_API_KEY")) self.api.auth_api_key(os.getenv("YETI_API_KEY")) - indicator = self.api.new_indicator( + self.api.new_indicator( { "name": "testSearch", "type": "regex", @@ -63,3 +63,11 @@ def test_search_indicators(self): result = self.api.search_indicators(name="testSear") self.assertEqual(len(result), 1, result) self.assertEqual(result[0]["name"], "testSearch") + + def test_find_indicator(self): + self.api.auth_api_key(os.getenv("YETI_API_KEY")) + self.api.auth_api_key(os.getenv("YETI_API_KEY")) + indicator = self.api.find_indicator(name="testSearch", type="regex") + + self.assertEqual(indicator["name"], "testSearch") + self.assertEqual(indicator["pattern"], "test[0-9]") From bcc045c237f25babcc484879f00905eb185867ca Mon Sep 17 00:00:00 2001 From: Thomas Chopitea Date: Fri, 2 May 2025 12:45:58 +0000 Subject: [PATCH 8/8] Update e2e --- tests/e2e.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/tests/e2e.py b/tests/e2e.py index 6516ca0..2f9a996 100644 --- a/tests/e2e.py +++ b/tests/e2e.py @@ -67,7 +67,17 @@ def test_search_indicators(self): def test_find_indicator(self): self.api.auth_api_key(os.getenv("YETI_API_KEY")) self.api.auth_api_key(os.getenv("YETI_API_KEY")) - indicator = self.api.find_indicator(name="testSearch", type="regex") + self.api.new_indicator( + { + "name": "testGet", + "type": "regex", + "description": "test", + "pattern": "test[0-9]", + "diamond": "victim", + } + ) + time.sleep(5) + indicator = self.api.find_indicator(name="testGet", type="regex") - self.assertEqual(indicator["name"], "testSearch") + self.assertEqual(indicator["name"], "testGet") self.assertEqual(indicator["pattern"], "test[0-9]")