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.
Taylor committed
Sep 3, 2024 at 11:58 UTC
cc3081076c7274b2ad651f5e17ecdf1ab4431afa
1 file changed
+31
-9
backend/app/agents/wazuh/services/vulnerabilities.py
+31
-9
@@ -68,7 +68,7 @@ async def collect_agent_vulnerabilities(agent_id: str, vulnerability_severity: s
68
)
69
if response["success"] is False:
70
raise HTTPException(status_code=500, detail=response["message"])
71
- agent_vulnerabilities.extend(response["data"])
71
+ agent_vulnerabilities.extend(response["data"].get("affected_items", []))
72
73
processed_vulnerabilities = process_agent_vulnerabilities(agent_vulnerabilities)
74
@@ -79,14 +79,41 @@ async def collect_agent_vulnerabilities(agent_id: str, vulnerability_severity: s
79
)
80
81
82
+# def process_agent_vulnerabilities(
83
+# agent_vulnerabilities: dict,
84
+# ) -> List[WazuhAgentVulnerabilities]:
85
+# """
86
+# Process agent vulnerabilities and return a list of WazuhAgentVulnerabilities objects.
87
+
88
+# Args:
89
+# agent_vulnerabilities (dict): A dictionary containing agent vulnerabilities data.
90
+
91
+# Returns:
92
+# List[WazuhAgentVulnerabilities]: A list of WazuhAgentVulnerabilities objects.
93
+
94
+# Raises:
95
+# HTTPException: If there is an error processing the agent vulnerabilities.
96
+# """
97
+# try:
98
+# vulnerabilities = agent_vulnerabilities.get("data", {}).get(
99
+# "affected_items",
100
+# [],
101
+# )
102
+# return [WazuhAgentVulnerabilities(**vuln) for vuln in vulnerabilities]
103
+# except Exception as e:
104
+# raise HTTPException(
105
+# status_code=500,
106
+# detail=f"Failed to process agent vulnerabilities: {e}",
107
+# )
108
+
109
def process_agent_vulnerabilities(
83
- agent_vulnerabilities: dict,
110
+ agent_vulnerabilities: List[dict],
111
) -> List[WazuhAgentVulnerabilities]:
112
"""
113
Process agent vulnerabilities and return a list of WazuhAgentVulnerabilities objects.
114
115
Args:
89
- agent_vulnerabilities (dict): A dictionary containing agent vulnerabilities data.
116
+ agent_vulnerabilities (List[dict]): A list of dictionaries containing agent vulnerabilities data.
117
118
Returns:
119
List[WazuhAgentVulnerabilities]: A list of WazuhAgentVulnerabilities objects.
@@ -95,18 +122,13 @@ def process_agent_vulnerabilities(
122
HTTPException: If there is an error processing the agent vulnerabilities.
123
"""
124
try:
98
- vulnerabilities = agent_vulnerabilities.get("data", {}).get(
99
- "affected_items",
100
- [],
101
- )
102
- return [WazuhAgentVulnerabilities(**vuln) for vuln in vulnerabilities]
125
+ return [WazuhAgentVulnerabilities(**vuln) for vuln in agent_vulnerabilities]
126
except Exception as e:
127
raise HTTPException(
128
status_code=500,
129
detail=f"Failed to process agent vulnerabilities: {e}",
130
)
131
109
-
132
async def collect_agent_vulnerabilities_new(agent_id: str, vulnerability_severity: str):
133
"""
134
Collects vulnerabilities for a specific agent from the Wazuh Indexer Index.