From 7998faa35450b0e0493b2c4324e3b438521c2c80 Mon Sep 17 00:00:00 2001 From: benyissa Date: Tue, 19 Dec 2023 12:18:54 +0100 Subject: [PATCH 1/4] add cidr limit. --- agent/virus_total_agent.py | 16 +++++++-- ostorlab.yaml | 2 +- tests/conftest.py | 40 ++++++++++++++++++++++ tests/virus_total_agent_test.py | 59 +++++++++++++++++++++++++++++++++ 4 files changed, 114 insertions(+), 3 deletions(-) diff --git a/agent/virus_total_agent.py b/agent/virus_total_agent.py index c2cf822..b1b64aa 100644 --- a/agent/virus_total_agent.py +++ b/agent/virus_total_agent.py @@ -27,6 +27,9 @@ logger = logging.getLogger(__name__) +IPV4_CIDR_LIMIT = 16 +IPV6_CIDR_LIMIT = 112 + class VirusTotalAgent( agent.Agent, agent_report_vulnerability_mixin.AgentReportVulnMixin @@ -106,10 +109,19 @@ def _prepare_targets(self, message: msg.Message) -> list[str]: from the config.""" if message.data.get("host") is not None: host = str(message.data.get("host")) - if message.data.get("mask") is None: + mask = message.data.get("mask") + if mask is None: ip_network = ipaddress.ip_network(host) else: - mask = message.data.get("mask") + version = message.data.get("version") + if version == 4 and int(mask) < IPV4_CIDR_LIMIT: + raise ValueError( + f"Subnet mask below {IPV4_CIDR_LIMIT} is not supported." + ) + if version == 6 and int(mask) < IPV6_CIDR_LIMIT: + raise ValueError( + f"Subnet mask below {IPV6_CIDR_LIMIT} is not supported." + ) ip_network = ipaddress.ip_network(f"{host}/{mask}", strict=False) return [str(h) for h in ip_network.hosts()] diff --git a/ostorlab.yaml b/ostorlab.yaml index 74ad695..59b0d63 100644 --- a/ostorlab.yaml +++ b/ostorlab.yaml @@ -1,6 +1,6 @@ kind: Agent name: virustotal -version: 0.2.0 +version: 0.2.1 image: images/logo.png description: | This repository is an implementation of the VirusTotal agent. diff --git a/tests/conftest.py b/tests/conftest.py index e4cc32f..a8a5662 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -129,3 +129,43 @@ def message_without_path() -> msg.Message: selector = "v3.asset.file" msg_data = {"content": file_content} return msg.Message.from_data(selector, data=msg_data) + + +@pytest.fixture() +def scan_message_ipv4_with_mask8() -> msg.Message: + """Creates a message of type v3.asset.ip.v4 to be used by the agent for testing purposes.""" + selector = "v3.asset.ip.v4" + msg_data = {"host": "192.168.1.17", "mask": "8", "version": 4} + return msg.Message.from_data(selector, data=msg_data) + + +@pytest.fixture() +def scan_message_ipv4_with_mask16() -> msg.Message: + """Creates a message of type v3.asset.ip.v4 to be used by the agent for testing purposes.""" + selector = "v3.asset.ip.v4" + msg_data = {"host": "192.168.1.17", "mask": "16", "version": 4} + return msg.Message.from_data(selector, data=msg_data) + + +@pytest.fixture() +def scan_message_ipv6_with_mask64() -> msg.Message: + """Creates a message of type v3.asset.ip.v6 to be used by the agent for testing purposes.""" + selector = "v3.asset.ip.v6" + msg_data = { + "host": "2001:db8:3333:4444:5555:6666:7777:8888", + "mask": "64", + "version": 6, + } + return msg.Message.from_data(selector, data=msg_data) + + +@pytest.fixture() +def scan_message_ipv6_with_mask112() -> msg.Message: + """Creates a message of type v3.asset.ip.v6 to be used by the agent for testing purposes.""" + selector = "v3.asset.ip.v6" + msg_data = { + "host": "2001:db8:3333:4444:5555:6666:7777:8888", + "mask": "112", + "version": 6, + } + return msg.Message.from_data(selector, data=msg_data) diff --git a/tests/virus_total_agent_test.py b/tests/virus_total_agent_test.py index e2891e8..29918c2 100644 --- a/tests/virus_total_agent_test.py +++ b/tests/virus_total_agent_test.py @@ -6,6 +6,7 @@ from ostorlab.agent.message import message as msg from pytest_mock import plugin +import pytest from agent import virus_total_agent from agent import virustotal @@ -334,3 +335,61 @@ def testVirusTotalAgent_whenFileHasNoPath_shouldReportWithHash( "Analysis of the target `44d88612fea8a8f36de82e1278abb02f`:\n|Package| Result | \n" "|-------|----------| \n|Bkav |_Safe_ | \n|Elastic|_Malicous_| \n" ) + + +def testPrepareTargets_whenIPv4AssetReachCIDRLimit_raiseValueError( + mocker: plugin.MockerFixture, + scan_message_ipv4_with_mask8: msg.Message, + virustotal_agent: virus_total_agent.VirusTotalAgent, +) -> None: + """Test the CIDR Limit in case IPV4 and the Limit is reached.""" + mocker.patch( + "agent.virustotal.scan_url_from_message", + return_value={}, + ) + + with pytest.raises(ValueError, match="Subnet mask below 16 is not supported."): + virustotal_agent.process(scan_message_ipv4_with_mask8) + + +def testPrepareTargets_whenIPv4AssetDoesNotReachCIDRLimit_doesNotRaiseValueError( + mocker: plugin.MockerFixture, + scan_message_ipv4_with_mask16: msg.Message, + virustotal_agent: virus_total_agent.VirusTotalAgent, +) -> None: + """Test the CIDR Limit in case IPV4 and the Limit is not reached.""" + mocker.patch( + "agent.virustotal.scan_url_from_message", + return_value={}, + ) + + virustotal_agent.process(scan_message_ipv4_with_mask16) + + +def testPrepareTargets_whenIPv6AssetReachCIDRLimit_raiseValueError( + mocker: plugin.MockerFixture, + scan_message_ipv6_with_mask64: msg.Message, + virustotal_agent: virus_total_agent.VirusTotalAgent, +) -> None: + """Test the CIDR Limit in case IPV6 and the Limit is reached.""" + mocker.patch( + "agent.virustotal.scan_url_from_message", + return_value={}, + ) + + with pytest.raises(ValueError, match="Subnet mask below 112 is not supported."): + virustotal_agent.process(scan_message_ipv6_with_mask64) + + +def testPrepareTargets_whenIPv6AssetDoesNotReachCIDRLimit_doesNotRaiseValueError( + mocker: plugin.MockerFixture, + scan_message_ipv6_with_mask112: msg.Message, + virustotal_agent: virus_total_agent.VirusTotalAgent, +) -> None: + """Test the CIDR Limit in case IPV6 and the Limit is not reached.""" + mocker.patch( + "agent.virustotal.scan_url_from_message", + return_value={}, + ) + + virustotal_agent.process(scan_message_ipv6_with_mask112) From eee2c4fa69cd93113a807a7300cdd75e2c565c73 Mon Sep 17 00:00:00 2001 From: benyissa Date: Tue, 19 Dec 2023 14:32:22 +0100 Subject: [PATCH 2/4] add cidr limit. --- agent/virus_total_agent.py | 6 ++++-- tests/virus_total_agent_test.py | 20 ++++---------------- 2 files changed, 8 insertions(+), 18 deletions(-) diff --git a/agent/virus_total_agent.py b/agent/virus_total_agent.py index b1b64aa..a55261c 100644 --- a/agent/virus_total_agent.py +++ b/agent/virus_total_agent.py @@ -114,11 +114,13 @@ def _prepare_targets(self, message: msg.Message) -> list[str]: ip_network = ipaddress.ip_network(host) else: version = message.data.get("version") - if version == 4 and int(mask) < IPV4_CIDR_LIMIT: + if version != 4 and version != 6: + raise ValueError(f"Incorrect ip version {version}") + elif version == 4 and int(mask) < IPV4_CIDR_LIMIT: raise ValueError( f"Subnet mask below {IPV4_CIDR_LIMIT} is not supported." ) - if version == 6 and int(mask) < IPV6_CIDR_LIMIT: + elif version == 6 and int(mask) < IPV6_CIDR_LIMIT: raise ValueError( f"Subnet mask below {IPV6_CIDR_LIMIT} is not supported." ) diff --git a/tests/virus_total_agent_test.py b/tests/virus_total_agent_test.py index 29918c2..5e41292 100644 --- a/tests/virus_total_agent_test.py +++ b/tests/virus_total_agent_test.py @@ -337,22 +337,16 @@ def testVirusTotalAgent_whenFileHasNoPath_shouldReportWithHash( ) -def testPrepareTargets_whenIPv4AssetReachCIDRLimit_raiseValueError( - mocker: plugin.MockerFixture, +def testVirusTotalAgent_whenIPv4AssetReachCIDRLimit_raiseValueError( scan_message_ipv4_with_mask8: msg.Message, virustotal_agent: virus_total_agent.VirusTotalAgent, ) -> None: """Test the CIDR Limit in case IPV4 and the Limit is reached.""" - mocker.patch( - "agent.virustotal.scan_url_from_message", - return_value={}, - ) - with pytest.raises(ValueError, match="Subnet mask below 16 is not supported."): virustotal_agent.process(scan_message_ipv4_with_mask8) -def testPrepareTargets_whenIPv4AssetDoesNotReachCIDRLimit_doesNotRaiseValueError( +def testVirusTotalAgent_whenIPv4AssetDoesNotReachCIDRLimit_doesNotRaiseValueError( mocker: plugin.MockerFixture, scan_message_ipv4_with_mask16: msg.Message, virustotal_agent: virus_total_agent.VirusTotalAgent, @@ -366,22 +360,16 @@ def testPrepareTargets_whenIPv4AssetDoesNotReachCIDRLimit_doesNotRaiseValueError virustotal_agent.process(scan_message_ipv4_with_mask16) -def testPrepareTargets_whenIPv6AssetReachCIDRLimit_raiseValueError( - mocker: plugin.MockerFixture, +def testVirusTotalAgent_whenIPv6AssetReachCIDRLimit_raiseValueError( scan_message_ipv6_with_mask64: msg.Message, virustotal_agent: virus_total_agent.VirusTotalAgent, ) -> None: """Test the CIDR Limit in case IPV6 and the Limit is reached.""" - mocker.patch( - "agent.virustotal.scan_url_from_message", - return_value={}, - ) - with pytest.raises(ValueError, match="Subnet mask below 112 is not supported."): virustotal_agent.process(scan_message_ipv6_with_mask64) -def testPrepareTargets_whenIPv6AssetDoesNotReachCIDRLimit_doesNotRaiseValueError( +def testVirusTotalAgent_whenIPv6AssetDoesNotReachCIDRLimit_doesNotRaiseValueError( mocker: plugin.MockerFixture, scan_message_ipv6_with_mask112: msg.Message, virustotal_agent: virus_total_agent.VirusTotalAgent, From 9d97c69364e39915d2b7d54309539abf5842cdc2 Mon Sep 17 00:00:00 2001 From: benyissa Date: Tue, 19 Dec 2023 14:34:30 +0100 Subject: [PATCH 3/4] add cidr limit. --- agent/virus_total_agent.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agent/virus_total_agent.py b/agent/virus_total_agent.py index a55261c..77d47b2 100644 --- a/agent/virus_total_agent.py +++ b/agent/virus_total_agent.py @@ -114,7 +114,7 @@ def _prepare_targets(self, message: msg.Message) -> list[str]: ip_network = ipaddress.ip_network(host) else: version = message.data.get("version") - if version != 4 and version != 6: + if version not in (4, 6): raise ValueError(f"Incorrect ip version {version}") elif version == 4 and int(mask) < IPV4_CIDR_LIMIT: raise ValueError( From 591e4876dac7b992fbe394a07c3f85123cee6cdb Mon Sep 17 00:00:00 2001 From: benyissa Date: Tue, 19 Dec 2023 14:49:34 +0100 Subject: [PATCH 4/4] add cidr limit. --- agent/virus_total_agent.py | 2 +- tests/conftest.py | 12 ++++++++++++ tests/virus_total_agent_test.py | 9 +++++++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/agent/virus_total_agent.py b/agent/virus_total_agent.py index 77d47b2..e2ba4cd 100644 --- a/agent/virus_total_agent.py +++ b/agent/virus_total_agent.py @@ -115,7 +115,7 @@ def _prepare_targets(self, message: msg.Message) -> list[str]: else: version = message.data.get("version") if version not in (4, 6): - raise ValueError(f"Incorrect ip version {version}") + raise ValueError(f"Incorrect ip version {version}.") elif version == 4 and int(mask) < IPV4_CIDR_LIMIT: raise ValueError( f"Subnet mask below {IPV4_CIDR_LIMIT} is not supported." diff --git a/tests/conftest.py b/tests/conftest.py index a8a5662..6833637 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -169,3 +169,15 @@ def scan_message_ipv6_with_mask112() -> msg.Message: "version": 6, } return msg.Message.from_data(selector, data=msg_data) + + +@pytest.fixture() +def scan_message_ipv_with_incorrect_version() -> msg.Message: + """Creates a message of type v3.asset.ip with an incorrect version.""" + selector = "v3.asset.ip" + msg_data = { + "host": "0.0.0.0", + "mask": "32", + "version": 5, + } + return msg.Message.from_data(selector, data=msg_data) diff --git a/tests/virus_total_agent_test.py b/tests/virus_total_agent_test.py index 5e41292..d07a80a 100644 --- a/tests/virus_total_agent_test.py +++ b/tests/virus_total_agent_test.py @@ -381,3 +381,12 @@ def testVirusTotalAgent_whenIPv6AssetDoesNotReachCIDRLimit_doesNotRaiseValueErro ) virustotal_agent.process(scan_message_ipv6_with_mask112) + + +def testVirusTotalAgent_whenIPAssetHasIncorrectVersion_raiseValueError( + scan_message_ipv_with_incorrect_version: msg.Message, + virustotal_agent: virus_total_agent.VirusTotalAgent, +) -> None: + """Test the CIDR Limit in case IP has incorrect version.""" + with pytest.raises(ValueError, match="Incorrect ip version 5."): + virustotal_agent.process(scan_message_ipv_with_incorrect_version)