Catches ASNParserError and logs

Ostorlab · Dec 2, 2024 · 0e83053 · 0e83053
1 parent 835d8f3
commit 0e83053
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 11 deletions.
diff --git a/agent/whois_ip_agent.py b/agent/whois_ip_agent.py
@@ -39,6 +39,7 @@
     reraise=True,
 )
 def _get_whois_record(host: str) -> dict[str, Any]:
+    logger.info("host: %s", host)
     lookup_rdap = ipwhois.IPWhois(host).lookup_rdap()
     return cast(dict[str, Any], lookup_rdap)
 
@@ -64,7 +65,7 @@ def process(self, message: m.Message) -> None:
         Returns:
             None
         """
-        logger.debug("processing message of selector %s", message.selector)
+        logger.info("processing message of selector %s", message.selector)
         if message.selector.startswith("v3.asset.domain_name.dns_record"):
             return self._process_dns_record(message)
         host = message.data.get("host")
@@ -153,14 +154,15 @@ def _process_ip(self, message: m.Message, host: str) -> None:
                     ipwhois.exceptions.IPDefinedError,
                     ipwhois.exceptions.ASNRegistryError,
                     ipwhois.exceptions.HTTPLookupError,
+                    ipwhois.exceptions.ASNParseError,
                 ):
                     # Case where of the loopback address.
-                    logger.warning(
+                    logger.error(
                         "Some data not found when agent_whois_ip_asset try to process IP %s",
                         address,
                     )
                 except exceptions.HTTPRateLimitError:
-                    logger.warning("Rate limit error for IP %s", address)
+                    logger.info("Rate limit error for IP %s", address)
         else:
             logger.info("target %s was processed before, exiting", network)
             return

diff --git a/tests/whois_ip_agent_test.py b/tests/whois_ip_agent_test.py
@@ -318,14 +318,12 @@ def testWhoisIP_withIPv4AndMaskButNoVersion_shouldHandleVersionCorrectly(
         data=message_data,
     )
 
-    with mock.patch.object(
-        test_agent, "_redis_client"
-    ) as mock_redis_client, mock.patch.object(
-        test_agent, "add_ip_network"
-    ) as mock_add_ip_network, mock.patch.object(
-        test_agent, "start", mock.MagicMock()
-    ), mock.patch.object(test_agent, "run", mock.MagicMock()), mock.patch(
-        "agent.whois_ip_agent.WhoisIPAgent.main", mock.MagicMock()
+    with (
+        mock.patch.object(test_agent, "_redis_client") as mock_redis_client,
+        mock.patch.object(test_agent, "add_ip_network") as mock_add_ip_network,
+        mock.patch.object(test_agent, "start", mock.MagicMock()),
+        mock.patch.object(test_agent, "run", mock.MagicMock()),
+        mock.patch("agent.whois_ip_agent.WhoisIPAgent.main", mock.MagicMock()),
     ):
         mock_redis_client.sismember.return_value = False
 
@@ -347,3 +345,24 @@ def testWhoisIP_whenInvalidIPAddressIsProvided_raisesValueError(
 
     with pytest.raises(ValueError, match="Invalid IP address: invalid_ip"):
         test_agent.process(ip_msg)
+
+
+def testWhoisIp_whenASNParseErrorOccure_logWithoutCrash(
+    test_agent: whois_ip_agent.WhoisIPAgent,
+    scan_message_ipv4: message.Message,
+    agent_persist_mock: dict[str | bytes, str | bytes],
+    mocker: plugin.MockerFixture,
+    agent_mock: List[message.Message],
+    caplog: pytest.LogCaptureFixture,
+):
+    """Test that ASNParseError is caught and handled gracefully."""
+    mocker.patch(
+        "ipwhois.IPWhois.lookup_rdap", side_effect=ipwhois.exceptions.ASNParseError
+    )
+
+    test_agent.process(scan_message_ipv4)
+
+    assert len(agent_mock) == 0
+    assert (
+        "Some data not found when agent_whois_ip_asset try to process IP" in caplog.text
+    )