From 9bf9180b203f8a7bf50834e5b1b36056e719f3d9 Mon Sep 17 00:00:00 2001
From: deadly-panda <abderrahim.haddadi@ostorlab.dev>
Date: Wed, 4 Dec 2024 18:02:17 +0100
Subject: [PATCH] Handle pysnmp_errors.PySnmpError gracefully.

---
 agent/exploits/cve_2024_40766.py      | 18 +++++++++++-------
 tests/exploits/cve_2024_40766_test.py | 16 ++++++++++++++++
 2 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/agent/exploits/cve_2024_40766.py b/agent/exploits/cve_2024_40766.py
index e814cbee..be6e0ae6 100644
--- a/agent/exploits/cve_2024_40766.py
+++ b/agent/exploits/cve_2024_40766.py
@@ -4,6 +4,7 @@
 import re
 
 from pysnmp import hlapi
+from pysnmp import error as pysnmp_errors
 from ostorlab.agent.kb import kb
 from ostorlab.agent.mixins import agent_report_vulnerability_mixin
 from agent import definitions
@@ -27,13 +28,16 @@
 
 def _get_sonicwall_version(host: str) -> str | None:
     version = None
-    iterator = hlapi.getCmd(
-        hlapi.SnmpEngine(),
-        hlapi.CommunityData("public", mpModel=1),
-        hlapi.UdpTransportTarget((host, 161), timeout=DEFAULT_TIMEOUT),
-        hlapi.ContextData(),
-        hlapi.ObjectType(hlapi.ObjectIdentity("1.3.6.1.2.1.1.1.0")),  # sysDescr OID
-    )
+    try:
+        iterator = hlapi.getCmd(
+            hlapi.SnmpEngine(),
+            hlapi.CommunityData("public", mpModel=1),
+            hlapi.UdpTransportTarget((host, 161), timeout=DEFAULT_TIMEOUT),
+            hlapi.ContextData(),
+            hlapi.ObjectType(hlapi.ObjectIdentity("1.3.6.1.2.1.1.1.0")),  # sysDescr OID
+        )
+    except pysnmp_errors.PySnmpError:
+        return None
     error_indication, error_status, error_index, var_binds = next(iterator)
 
     if error_indication is True:
diff --git a/tests/exploits/cve_2024_40766_test.py b/tests/exploits/cve_2024_40766_test.py
index 842bfcbf..03461941 100644
--- a/tests/exploits/cve_2024_40766_test.py
+++ b/tests/exploits/cve_2024_40766_test.py
@@ -3,6 +3,7 @@
 from unittest import mock
 
 from pytest_mock import plugin
+from pysnmp import error as pysnmp_errors
 
 from agent import definitions
 from agent.exploits import cve_2024_40766
@@ -112,3 +113,18 @@ def testCVE202440766_whenSNMPError_handleGracefully(
     assert accept is False
     assert len(vulnerabilities) == 0
     mock_logging.assert_called_with("error_indication: %s", True)
+
+
+def testCVE202440766_whenPySnmpErrorRaise_handleGracefully(
+    mocker: plugin.MockerFixture,
+) -> None:
+    """Ensure CVE-2024-40766 exploit handles PySnmpError raised by hlapi.getCmd."""
+    mocker.patch("pysnmp.hlapi.getCmd", side_effect=pysnmp_errors.PySnmpError)
+    exploit_instance = cve_2024_40766.CVE202440766Exploit()
+    target = definitions.Target("udp", "192.168.1.1", 161)
+
+    accept = exploit_instance.accept(target)
+    vulnerabilities = exploit_instance.check(target)
+
+    assert accept is False
+    assert len(vulnerabilities) == 0