add cidr limit.

Ostorlab · Dec 19, 2023 · 7ac039c · 7ac039c
1 parent b904ea7
commit 7ac039c
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 6 deletions.
diff --git a/agent/whois_ip_agent.py b/agent/whois_ip_agent.py
@@ -121,11 +121,13 @@ def _process_ip(self, message: m.Message) -> None:
             network = ipaddress.ip_network(f"{host}")
         else:
             version = message.data.get("version")
-            if version == 4 and int(mask) < IPV4_CIDR_LIMIT:
+            if version not in (4, 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/conftest.py b/tests/conftest.py
@@ -149,3 +149,15 @@ def scan_message_ipv6_with_mask112() -> message.Message:
         "version": 6,
     }
     return message.Message.from_data(selector, data=msg_data)
+
+
+@pytest.fixture()
+def scan_message_ipv_with_incorrect_version() -> message.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 message.Message.from_data(selector, data=msg_data)
diff --git a/tests/whois_ip_agent_test.py b/tests/whois_ip_agent_test.py
@@ -228,7 +228,7 @@ def testAgentWhoisIP_whenRDAPIsDown_shouldRetry(
     assert mock_request.call_count == 2
 
 
-def testPrepareTargets_whenIPv4AssetReachCIDRLimit_raiseValueError(
+def testWhoisIP_whenIPv4AssetReachCIDRLimit_raiseValueError(
     test_agent: whois_ip_agent.WhoisIPAgent,
     mocker: plugin.MockerFixture,
     scan_message_ipv4_with_mask8: message.Message,
@@ -243,7 +243,7 @@ def testPrepareTargets_whenIPv4AssetReachCIDRLimit_raiseValueError(
         test_agent.process(scan_message_ipv4_with_mask8)
 
 
-def testPrepareTargets_whenIPv4AssetDoesNotReachCIDRLimit_doesNotRaiseValueError(
+def testWhoisIP_whenIPv4AssetDoesNotReachCIDRLimit_doesNotRaiseValueError(
     test_agent: whois_ip_agent.WhoisIPAgent,
     mocker: plugin.MockerFixture,
     scan_message_ipv4_with_mask16: message.Message,
@@ -257,7 +257,7 @@ def testPrepareTargets_whenIPv4AssetDoesNotReachCIDRLimit_doesNotRaiseValueError
     test_agent.process(scan_message_ipv4_with_mask16)
 
 
-def testPrepareTargets_whenIPv6AssetReachCIDRLimit_raiseValueError(
+def testWhoisIP_whenIPv6AssetReachCIDRLimit_raiseValueError(
     test_agent: whois_ip_agent.WhoisIPAgent,
     mocker: plugin.MockerFixture,
     scan_message_ipv6_with_mask64: message.Message,
@@ -272,7 +272,7 @@ def testPrepareTargets_whenIPv6AssetReachCIDRLimit_raiseValueError(
         test_agent.process(scan_message_ipv6_with_mask64)
 
 
-def testPrepareTargets_whenIPv6AssetDoesNotReachCIDRLimit_doesNotRaiseValueError(
+def testWhoisIP_whenIPv6AssetDoesNotReachCIDRLimit_doesNotRaiseValueError(
     test_agent: whois_ip_agent.WhoisIPAgent,
     mocker: plugin.MockerFixture,
     scan_message_ipv6_with_mask112: message.Message,
@@ -284,3 +284,12 @@ def testPrepareTargets_whenIPv6AssetDoesNotReachCIDRLimit_doesNotRaiseValueError
     )
 
     test_agent.process(scan_message_ipv6_with_mask112)
+
+
+def testWhoisIP_whenIPAssetHasIncorrectVersion_raiseValueError(
+    test_agent: whois_ip_agent.WhoisIPAgent,
+    scan_message_ipv_with_incorrect_version: message.Message,
+) -> None:
+    """Test the CIDR Limit in case IP has incorrect version."""
+    with pytest.raises(ValueError, match="Incorrect ip version 5."):
+        test_agent.process(scan_message_ipv_with_incorrect_version)