From 72c87e4337813945cb0a5b578e82f4581e0b5b44 Mon Sep 17 00:00:00 2001 From: ohachimOs Date: Thu, 5 Dec 2024 13:54:53 +0100 Subject: [PATCH 1/8] Mock operations that take too much time --- tests/asteroid_agent_test.py | 40 +++++++++++++++++++++++++++--------- 1 file changed, 30 insertions(+), 10 deletions(-) diff --git a/tests/asteroid_agent_test.py b/tests/asteroid_agent_test.py index 861001ea..1373c965 100644 --- a/tests/asteroid_agent_test.py +++ b/tests/asteroid_agent_test.py @@ -1,11 +1,14 @@ """Unit tests for AsteroidAgent.""" -from typing import Type, Iterator +from typing import Callable, Iterator, Type +import requests +import requests_mock from ostorlab.agent.message import message as m +from pytest_mock import plugin +from requests_mock.adapter import ANY -from agent import asteroid_agent -from agent import definitions +from agent import asteroid_agent, definitions def testAsteroidAgent_whenExploitCheckDetectVulnz_EmitsVulnerabilityReport( @@ -23,18 +26,35 @@ def testAsteroidAgent_whenExploitCheckDetectVulnz_EmitsVulnerabilityReport( def testAsteroidAgent_whenTooManyRedirects_doesNotCrash( - exploit_instance_with_report: Iterator[Type[definitions.Exploit]], asteroid_agent_instance: asteroid_agent.AsteroidAgent, agent_mock: list[m.Message], + mocker: plugin.MockerFixture, + requests_mock: requests_mock.Mocker, ) -> None: """Ensure that the agent does not crash when there are too many redirects.""" + + def response_callback(request: requests.Request, context: Callable) -> str: + context.headers = {"Location": request.url} + context.status_code = 302 + return "" + + requests_mock.register_uri( + ANY, + ANY, + text=response_callback, + ) + + mock_var_bind = mocker.MagicMock() + mock_var_bind.__getitem__.return_value.prettyPrint.return_value = ( + "ArubaOS (MODEL: 7005), Version 8.5.0.0" + ) + mock_iterator = mocker.MagicMock() + mock_iterator.__next__.return_value = (None, None, None, [mock_var_bind]) + mocker.patch("pysnmp.hlapi.getCmd", return_value=mock_iterator) + msg = m.Message( selector="v3.asset.link", - data={"url": "https://expediaagents.com", "method": "GET"}, - raw=b"\n\x19https://expediaagents.com\x12\x03GET", + data={"url": "https://example.com", "method": "GET"}, + raw=b"\n\x19https://example.com\x12\x03GET", ) - asteroid_agent_instance.process(msg) - - assert len(agent_mock) == 1 - assert agent_mock[0].selector == "v3.report.vulnerability" From 43d185fdf94a2cc6f755b988abe8c79d495ab246 Mon Sep 17 00:00:00 2001 From: ohachimOs Date: Thu, 5 Dec 2024 14:01:32 +0100 Subject: [PATCH 2/8] Add assert --- tests/asteroid_agent_test.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/asteroid_agent_test.py b/tests/asteroid_agent_test.py index 1373c965..06387393 100644 --- a/tests/asteroid_agent_test.py +++ b/tests/asteroid_agent_test.py @@ -3,6 +3,7 @@ from typing import Callable, Iterator, Type import requests +import agent import requests_mock from ostorlab.agent.message import message as m from pytest_mock import plugin @@ -58,3 +59,5 @@ def response_callback(request: requests.Request, context: Callable) -> str: raw=b"\n\x19https://example.com\x12\x03GET", ) asteroid_agent_instance.process(msg) + + assert len(agent_mock) == 1 From d2237b1a6ab07010c3c117630fd9e1d627973c89 Mon Sep 17 00:00:00 2001 From: ohachimOs Date: Thu, 5 Dec 2024 14:02:17 +0100 Subject: [PATCH 3/8] Remove import --- tests/asteroid_agent_test.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/asteroid_agent_test.py b/tests/asteroid_agent_test.py index 06387393..300a9267 100644 --- a/tests/asteroid_agent_test.py +++ b/tests/asteroid_agent_test.py @@ -3,7 +3,6 @@ from typing import Callable, Iterator, Type import requests -import agent import requests_mock from ostorlab.agent.message import message as m from pytest_mock import plugin From 8ceb92a2c32b5c6e8917ddc76c1ae98623532a7b Mon Sep 17 00:00:00 2001 From: ohachimOs Date: Thu, 5 Dec 2024 14:05:15 +0100 Subject: [PATCH 4/8] Fix mypy --- tests/asteroid_agent_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/asteroid_agent_test.py b/tests/asteroid_agent_test.py index 300a9267..7beb6c09 100644 --- a/tests/asteroid_agent_test.py +++ b/tests/asteroid_agent_test.py @@ -1,6 +1,6 @@ """Unit tests for AsteroidAgent.""" -from typing import Callable, Iterator, Type +from typing import Any, Iterator, Type import requests import requests_mock @@ -33,7 +33,7 @@ def testAsteroidAgent_whenTooManyRedirects_doesNotCrash( ) -> None: """Ensure that the agent does not crash when there are too many redirects.""" - def response_callback(request: requests.Request, context: Callable) -> str: + def response_callback(request: Any, context: Any) -> str: context.headers = {"Location": request.url} context.status_code = 302 return "" From bfb8d2aa08aa635d812de24d53c84517730a1a29 Mon Sep 17 00:00:00 2001 From: ohachimOs Date: Thu, 5 Dec 2024 14:12:43 +0100 Subject: [PATCH 5/8] ruff --- tests/asteroid_agent_test.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/asteroid_agent_test.py b/tests/asteroid_agent_test.py index 7beb6c09..5d2a3ba8 100644 --- a/tests/asteroid_agent_test.py +++ b/tests/asteroid_agent_test.py @@ -2,7 +2,6 @@ from typing import Any, Iterator, Type -import requests import requests_mock from ostorlab.agent.message import message as m from pytest_mock import plugin From f871b873e8ac5c4a7c65ca2b8e12a25089c2ba05 Mon Sep 17 00:00:00 2001 From: ohachimOs Date: Thu, 5 Dec 2024 14:14:35 +0100 Subject: [PATCH 6/8] Adds selector assertion back --- tests/asteroid_agent_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/asteroid_agent_test.py b/tests/asteroid_agent_test.py index 5d2a3ba8..cbfa12e1 100644 --- a/tests/asteroid_agent_test.py +++ b/tests/asteroid_agent_test.py @@ -59,3 +59,4 @@ def response_callback(request: Any, context: Any) -> str: asteroid_agent_instance.process(msg) assert len(agent_mock) == 1 + assert agent_mock[0].selector == "v3.report.vulnerability" From b34494807a74ded2a0d303991f6c43015c4603ea Mon Sep 17 00:00:00 2001 From: ohachimOs Date: Thu, 5 Dec 2024 15:18:51 +0100 Subject: [PATCH 7/8] Mocking to unittest --- tests/asteroid_agent_test.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/tests/asteroid_agent_test.py b/tests/asteroid_agent_test.py index cbfa12e1..2a1b6f87 100644 --- a/tests/asteroid_agent_test.py +++ b/tests/asteroid_agent_test.py @@ -15,12 +15,23 @@ def testAsteroidAgent_whenExploitCheckDetectVulnz_EmitsVulnerabilityReport( asteroid_agent_instance: asteroid_agent.AsteroidAgent, agent_mock: list[m.Message], scan_message_domain_name: m.Message, + mocker: plugin.MockerFixture, ) -> None: """Unit test for agent AsteroidAgent exploits check. case Exploit emits vulnerability report""" + mock_var_bind = mocker.MagicMock() + mock_var_bind.__getitem__.return_value.prettyPrint.return_value = ( + "ArubaOS (MODEL: 7005), Version 8.5.0.0" + ) + mock_iterator = mocker.MagicMock() + mock_iterator.__next__.return_value = (None, None, None, [mock_var_bind]) + mocker.patch("pysnmp.hlapi.getCmd", return_value=mock_iterator) + + requests_mock.register_uri(ANY, ANY, status_code=404, text="") + asteroid_agent_instance.process(scan_message_domain_name) - assert len(agent_mock) == 1 + assert len(agent_mock) > 0 assert agent_mock[0].selector == "v3.report.vulnerability" From 7f165617ba370dc307723cb28a0eb0cfbe4a6911 Mon Sep 17 00:00:00 2001 From: ohachimOs Date: Thu, 5 Dec 2024 15:22:25 +0100 Subject: [PATCH 8/8] Adds fixture --- tests/asteroid_agent_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/asteroid_agent_test.py b/tests/asteroid_agent_test.py index 2a1b6f87..5563b5c6 100644 --- a/tests/asteroid_agent_test.py +++ b/tests/asteroid_agent_test.py @@ -16,6 +16,7 @@ def testAsteroidAgent_whenExploitCheckDetectVulnz_EmitsVulnerabilityReport( agent_mock: list[m.Message], scan_message_domain_name: m.Message, mocker: plugin.MockerFixture, + requests_mock: requests_mock.Mocker, ) -> None: """Unit test for agent AsteroidAgent exploits check. case Exploit emits vulnerability report"""