From cc3081076c7274b2ad651f5e17ecdf1ab4431afa Mon Sep 17 00:00:00 2001
From: Taylor <taylor.walton@socfortress.co>
Date: Tue, 3 Sep 2024 11:58:53 -0500
Subject: [PATCH] feat: Update `collect_agent_vulnerabilities` to support
 multiple severity levels

The code changes in this commit modify the `collect_agent_vulnerabilities` function in the `vulnerabilities.py` file to support collecting agent vulnerabilities for multiple severity levels. Previously, the function only collected vulnerabilities for a single severity level. Now, it can collect vulnerabilities for all severity levels or a specific severity level specified by the user.
---
 .../agents/wazuh/services/vulnerabilities.py  | 40 ++++++++++++++-----
 1 file changed, 31 insertions(+), 9 deletions(-)

diff --git a/backend/app/agents/wazuh/services/vulnerabilities.py b/backend/app/agents/wazuh/services/vulnerabilities.py
index 4f220cd8..9a786a51 100644
--- a/backend/app/agents/wazuh/services/vulnerabilities.py
+++ b/backend/app/agents/wazuh/services/vulnerabilities.py
@@ -68,7 +68,7 @@ async def collect_agent_vulnerabilities(agent_id: str, vulnerability_severity: s
         )
         if response["success"] is False:
             raise HTTPException(status_code=500, detail=response["message"])
-        agent_vulnerabilities.extend(response["data"])
+        agent_vulnerabilities.extend(response["data"].get("affected_items", []))
 
     processed_vulnerabilities = process_agent_vulnerabilities(agent_vulnerabilities)
 
@@ -79,14 +79,41 @@ async def collect_agent_vulnerabilities(agent_id: str, vulnerability_severity: s
     )
 
 
+# def process_agent_vulnerabilities(
+#     agent_vulnerabilities: dict,
+# ) -> List[WazuhAgentVulnerabilities]:
+#     """
+#     Process agent vulnerabilities and return a list of WazuhAgentVulnerabilities objects.
+
+#     Args:
+#         agent_vulnerabilities (dict): A dictionary containing agent vulnerabilities data.
+
+#     Returns:
+#         List[WazuhAgentVulnerabilities]: A list of WazuhAgentVulnerabilities objects.
+
+#     Raises:
+#         HTTPException: If there is an error processing the agent vulnerabilities.
+#     """
+#     try:
+#         vulnerabilities = agent_vulnerabilities.get("data", {}).get(
+#             "affected_items",
+#             [],
+#         )
+#         return [WazuhAgentVulnerabilities(**vuln) for vuln in vulnerabilities]
+#     except Exception as e:
+#         raise HTTPException(
+#             status_code=500,
+#             detail=f"Failed to process agent vulnerabilities: {e}",
+#         )
+
 def process_agent_vulnerabilities(
-    agent_vulnerabilities: dict,
+    agent_vulnerabilities: List[dict],
 ) -> List[WazuhAgentVulnerabilities]:
     """
     Process agent vulnerabilities and return a list of WazuhAgentVulnerabilities objects.
 
     Args:
-        agent_vulnerabilities (dict): A dictionary containing agent vulnerabilities data.
+        agent_vulnerabilities (List[dict]): A list of dictionaries containing agent vulnerabilities data.
 
     Returns:
         List[WazuhAgentVulnerabilities]: A list of WazuhAgentVulnerabilities objects.
@@ -95,18 +122,13 @@ def process_agent_vulnerabilities(
         HTTPException: If there is an error processing the agent vulnerabilities.
     """
     try:
-        vulnerabilities = agent_vulnerabilities.get("data", {}).get(
-            "affected_items",
-            [],
-        )
-        return [WazuhAgentVulnerabilities(**vuln) for vuln in vulnerabilities]
+        return [WazuhAgentVulnerabilities(**vuln) for vuln in agent_vulnerabilities]
     except Exception as e:
         raise HTTPException(
             status_code=500,
             detail=f"Failed to process agent vulnerabilities: {e}",
         )
 
-
 async def collect_agent_vulnerabilities_new(agent_id: str, vulnerability_severity: str):
     """
     Collects vulnerabilities for a specific agent from the Wazuh Indexer Index.