Improve unit tests coverage in findings models

pablosnt · Dec 9, 2023 · b34ac96 · b34ac96
1 parent 76dcebb
commit b34ac96
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 11 deletions.
diff --git a/src/backend/findings/models.py b/src/backend/findings/models.py
@@ -30,13 +30,11 @@ class OSINT(Finding):
     def parse(
         self, target: Target = None, accumulated: Dict[str, Any] = {}
     ) -> Dict[str, Any]:
-        if self.data_type in [OSINTDataType.IP, OSINTDataType.DOMAIN]:
-            return {
-                InputKeyword.TARGET.name.lower(): self.data,
-                InputKeyword.HOST.name.lower(): self.data,
-                InputKeyword.URL.name.lower(): self._get_url(self.data),
-            }
-        return {}
+        return {
+            InputKeyword.TARGET.name.lower(): self.data,
+            InputKeyword.HOST.name.lower(): self.data,
+            InputKeyword.URL.name.lower(): self._get_url(self.data),
+        } if self.data_type in [OSINTDataType.IP, OSINTDataType.DOMAIN] else {}
 
     def defect_dojo(self) -> Dict[str, Any]:
         return {

diff --git a/src/backend/tests/executors/test_base.py b/src/backend/tests/executors/test_base.py
@@ -67,14 +67,14 @@ def test_get_arguments_only_findings(self) -> None:
     @mock.patch("framework.models.BaseInput._get_url", get_url)
     def test_get_arguments_only_required_findings(self) -> None:
         self._success_get_arguments(
-            "-p http://10.10.10.10:80/ -p 80 -p WordPress -p CVE-2023-1111",
+            "-p 10.10.10.10 -p http://10.10.10.10:80/ -p 80 -p WordPress -p CVE-2023-1111",
             [self.host, self.port, self.technology, self.vulnerability],
         )
 
     @mock.patch("framework.models.BaseInput._get_url", get_url)
     def test_get_arguments_multiple_ports(self) -> None:
         self._success_get_arguments(
-            "-p http://10.10.10.10:80/ -p 80,443 -p WordPress -p CVE-2023-1111",
+            "-p 10.10.10.10 -p http://10.10.10.10:80/ -p 80,443 -p WordPress -p CVE-2023-1111",
             [
                 self.host,
                 self.port,
@@ -90,7 +90,7 @@ def _test_get_arguments_no_findings(self) -> None:
         self.target.target = "10.10.10.12"
         self.target.save(update_fields=["target"])
         self._success_get_arguments(
-            f"-p http://10.10.10.12:80/login.php -p 80 -p /login.php -p Joomla -p CVE-2023-2222 -p {base64.b64encode('root:root'.encode()).decode() if self.authentication.type == AuthenticationType.BASIC else 'root'} -p {self.wordlist.path}",
+            f"-p 10.10.10.10 -p http://10.10.10.12:80/login.php -p 80 -p /login.php -p Joomla -p CVE-2023-2222 -p {base64.b64encode('root:root'.encode()).decode() if self.authentication.type == AuthenticationType.BASIC else 'root'} -p {self.wordlist.path}",
             [],
             [self.target_port],
             [self.input_vulnerability],

diff --git a/src/backend/tests/framework.py b/src/backend/tests/framework.py
@@ -118,7 +118,7 @@ def _setup_fake_tool(self) -> None:
             default=True,
         )
         for value, required, multiple, input_type_names in [
-            ("host", False, False, [InputTypeName.OSINT]),
+            ("host", False, False, [InputTypeName.OSINT, InputTypeName.HOST]),
             (
                 "url",
                 True,