Add network_device_common class + fix consistency

Signed-off-by: bretfourbe <[email protected]>
bretfourbe · Apr 10, 2024 · f0c605d · f0c605d
1 parent dec4f26
commit f0c605d
Show file tree

Hide file tree

Showing 6 changed files with 40 additions and 29 deletions.
diff --git a/tests/attack/test_mod_network_device.py b/tests/attack/test_mod_network_device.py
@@ -82,8 +82,9 @@ async def test_ubika_without_version():
 
         assert persister.add_payload.call_count == 1
         assert persister.add_payload.call_args_list[0][1]["info"] == (
-            '{"name": "UBIKA WAAP", "version": "", "categories": ["Network Equipment"], "groups": ["Content"]}'
+            '{"name": "UBIKA WAAP", "versions": [], "categories": ["Network Equipment"], "groups": ["Content"]}'
         )
+        assert persister.add_payload.call_args_list[0][1]["module"] == "network_device"
 
 @pytest.mark.asyncio
 @respx.mock
@@ -133,8 +134,9 @@ async def test_ubika_with_version():
         assert persister.add_payload.call_count == 1
         assert persister.add_payload.call_args_list[0][1]["category"] == "Fingerprint web technology"
         assert persister.add_payload.call_args_list[0][1]["info"] == (
-            '{"name": "UBIKA WAAP", "version": "6.5.6", "categories": ["Network Equipment"], "groups": ["Content"]}'
+            '{"name": "UBIKA WAAP", "versions": ["6.5.6"], "categories": ["Network Equipment"], "groups": ["Content"]}'
         )
+        assert persister.add_payload.call_args_list[0][1]["module"] == "network_device"
 
 
 @pytest.mark.asyncio
@@ -174,8 +176,9 @@ async def test_detect_fortimanager():
 
         assert persister.add_payload.call_count == 1
         assert persister.add_payload.call_args_list[0][1]["info"] == (
-            '{"name": "FortiManager", "version": "", "categories": ["Network Equipment"], "groups": ["Content"]}'
+            '{"name": "FortiManager", "versions": [], "categories": ["Network Equipment"], "groups": ["Content"]}'
         )
+        assert persister.add_payload.call_args_list[0][1]["module"] == "network_device"
 
 
 @pytest.mark.asyncio
@@ -214,8 +217,9 @@ async def test_detect_ssl_vpn():
 
         assert persister.add_payload.call_count == 1
         assert persister.add_payload.call_args_list[0][1]["info"] == (
-            '{"name": "Fortinet SSL-VPN", "version": "", "categories": ["Network Equipment"], "groups": ["Content"]}'
+            '{"name": "Fortinet SSL-VPN", "versions": [], "categories": ["Network Equipment"], "groups": ["Content"]}'
         )
+        assert persister.add_payload.call_args_list[0][1]["module"] == "network_device"
 
 
 @pytest.mark.asyncio
@@ -253,8 +257,9 @@ async def test_detect_fortinet():
 
         assert persister.add_payload.call_count == 1
         assert persister.add_payload.call_args_list[0][1]["info"] == (
-            '{"name": "Fortinet", "version": "", "categories": ["Network Equipment"], "groups": ["Content"]}'
+            '{"name": "Fortinet", "versions": [], "categories": ["Network Equipment"], "groups": ["Content"]}'
         )
+        assert persister.add_payload.call_args_list[0][1]["module"] == "network_device"
 
 
 @pytest.mark.asyncio
@@ -292,8 +297,9 @@ async def test_detect_fortiportal_from_title():
 
         assert persister.add_payload.call_count == 1
         assert persister.add_payload.call_args_list[0][1]["info"] == (
-            '{"name": "FortiPortal", "version": "", "categories": ["Network Equipment"], "groups": ["Content"]}'
+            '{"name": "FortiPortal", "versions": [], "categories": ["Network Equipment"], "groups": ["Content"]}'
         )
+        assert persister.add_payload.call_args_list[0][1]["module"] == "network_device"
 
 
 @pytest.mark.asyncio
@@ -332,8 +338,9 @@ async def test_detect_fortimail():
 
         assert persister.add_payload.call_count == 1
         assert persister.add_payload.call_args_list[0][1]["info"] == (
-            '{"name": "FortiMail", "version": "", "categories": ["Network Equipment"], "groups": ["Content"]}'
+            '{"name": "FortiMail", "versions": [], "categories": ["Network Equipment"], "groups": ["Content"]}'
         )
+        assert persister.add_payload.call_args_list[0][1]["module"] == "network_device"
 
 
 @pytest.mark.asyncio
@@ -412,8 +419,9 @@ async def test_detect_harbor_with_version():
 
         assert persister.add_payload.call_count == 1
         assert persister.add_payload.call_args_list[0][1]["info"] == (
-            '{"name": "Harbor", "version": "v2.10", "categories": ["Network Equipment"], "groups": ["Content"]}'
+            '{"name": "Harbor", "versions": ["v2.10"], "categories": ["Network Equipment"], "groups": ["Content"]}'
         )
+        assert persister.add_payload.call_args_list[0][1]["module"] == "network_device"
 
 @pytest.mark.asyncio
 @respx.mock
@@ -458,8 +466,9 @@ async def test_detect_harbor_without_version():
 
         assert persister.add_payload.call_count == 1
         assert persister.add_payload.call_args_list[0][1]["info"] == (
-            '{"name": "Harbor", "version": "", "categories": ["Network Equipment"], "groups": ["Content"]}'
+            '{"name": "Harbor", "versions": [], "categories": ["Network Equipment"], "groups": ["Content"]}'
         )
+        assert persister.add_payload.call_args_list[0][1]["module"] == "network_device"
 
 
 @pytest.mark.asyncio
@@ -499,8 +508,9 @@ async def test_detect_harbor_with_json_error():
 
         assert persister.add_payload.call_count == 1
         assert persister.add_payload.call_args_list[0][1]["info"] == (
-            '{"name": "Harbor", "version": "", "categories": ["Network Equipment"], "groups": ["Content"]}'
+            '{"name": "Harbor", "versions": [], "categories": ["Network Equipment"], "groups": ["Content"]}'
         )
+        assert persister.add_payload.call_args_list[0][1]["module"] == "network_device"
 
 
 @pytest.mark.asyncio

diff --git a/wapitiCore/attack/mod_network_device.py b/wapitiCore/attack/mod_network_device.py
@@ -9,10 +9,6 @@
 from wapitiCore.net.response import Response
 
 
-MSG_TECHNO_VERSIONED = "{0} {1} detected"
-MSG_NO_UBIKA = "No UBIKA Detected"
-
-
 class ModuleNetworkDevice(Attack):
     """Base class for detecting version."""
     name = "network_device"

diff --git a/wapitiCore/attack/network_devices/mod_forti.py b/wapitiCore/attack/network_devices/mod_forti.py
@@ -6,17 +6,17 @@
 from bs4 import BeautifulSoup
 from httpx import RequestError
 
-from wapitiCore.attack.attack import Attack
+from wapitiCore.attack.network_devices.network_device_common import NetworkDeviceCommon, MSG_TECHNO_VERSIONED
 from wapitiCore.net import Request
 from wapitiCore.net.response import Response
 from wapitiCore.definitions.fingerprint import NAME as TECHNO_DETECTED, WSTG_CODE
 from wapitiCore.main.log import log_blue, logging
 
 MSG_NO_FORTI = "No Forti Product Detected"
-MSG_FORTI_DETECTED = "{0}{1} Detected !"
 
 
-class ModuleForti(Attack):
+class ModuleForti(NetworkDeviceCommon):
+    """Detect Forti."""
     device_name = "Fortinet"
     version = ""
     fortinet_pattern = re.compile(r'Forti\w+')
@@ -129,12 +129,12 @@ async def attack(self, request: Request, response: Optional[Response] = None):
 
                 forti_detected = {
                     "name": self.device_name,
-                    "version": "",
+                    "versions": [],
                     "categories": ["Network Equipment"],
                     "groups": ["Content"]
                 }
                 log_blue(
-                    MSG_FORTI_DETECTED,
+                    MSG_TECHNO_VERSIONED,
                     self.device_name,
                     self.version
                 )

diff --git a/wapitiCore/attack/network_devices/mod_harbor.py b/wapitiCore/attack/network_devices/mod_harbor.py
@@ -5,17 +5,16 @@
 
 from httpx import RequestError
 
-from wapitiCore.attack.attack import Attack
+from wapitiCore.attack.network_devices.network_device_common import NetworkDeviceCommon, MSG_TECHNO_VERSIONED
 from wapitiCore.net import Request
 from wapitiCore.net.response import Response
 from wapitiCore.definitions.fingerprint import NAME as TECHNO_DETECTED, WSTG_CODE
 from wapitiCore.main.log import log_blue, logging
 
 MSG_NO_HARBOR = "No Harbor Product Detected"
-MSG_HARBOR_DETECTED = "{0} {1} Detected !"
 
 
-class ModuleHarbor(Attack):
+class ModuleHarbor(NetworkDeviceCommon):
     """Detect Harbor."""
 
     device_name = "Harbor"
@@ -61,12 +60,12 @@ async def attack(self, request: Request, response: Optional[Response] = None):
             if await self.check_harbor(request_to_root.url):
                 harbor_detected = {
                     "name": self.device_name,
-                    "version": self.version,
+                    "versions": [self.version] if self.version else [],
                     "categories": ["Network Equipment"],
                     "groups": ["Content"]
                 }
                 log_blue(
-                    MSG_HARBOR_DETECTED,
+                    MSG_TECHNO_VERSIONED,
                     self.device_name,
                     self.version
                 )

diff --git a/wapitiCore/attack/network_devices/mod_ubika.py b/wapitiCore/attack/network_devices/mod_ubika.py
@@ -6,18 +6,17 @@
 from bs4 import BeautifulSoup
 from httpx import RequestError
 
-from wapitiCore.attack.attack import Attack
+from wapitiCore.attack.network_devices.network_device_common import NetworkDeviceCommon, MSG_TECHNO_VERSIONED
 from wapitiCore.net import Request
 from wapitiCore.net.response import Response
 from wapitiCore.definitions.fingerprint import NAME as TECHNO_DETECTED, WSTG_CODE
 from wapitiCore.main.log import log_blue, logging
 
-MSG_TECHNO_VERSIONED = "{0} {1} detected"
 MSG_NO_UBIKA = "No UBIKA Detected"
 
 
-class ModuleUbika(Attack):
-    """Base class for detecting version."""
+class ModuleUbika(NetworkDeviceCommon):
+    """Detect Ubika."""
     version = ""
 
     async def check_ubika(self, url):
@@ -63,7 +62,7 @@ async def attack(self, request: Request, response: Optional[Response] = None):
 
                 ubika_detected = {
                     "name": "UBIKA WAAP",
-                    "version": self.version,
+                    "versions": [self.version] if self.version else [],
                     "categories": ["Network Equipment"],
                     "groups": ["Content"]
                 }

diff --git a/wapitiCore/attack/network_devices/network_device_common.py b/wapitiCore/attack/network_devices/network_device_common.py
@@ -0,0 +1,7 @@
+from wapitiCore.attack.attack import Attack
+
+MSG_TECHNO_VERSIONED = "{0} {1} detected"
+
+class NetworkDeviceCommon(Attack):
+    """Base class for detecting version."""
+    name = "network_device"