Skip to content

Commit

Permalink
feat: Update collect_agent_vulnerabilities to support multiple seve…
Browse files Browse the repository at this point in the history
…rity 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.
  • Loading branch information
taylorwalton committed Sep 3, 2024
1 parent c8eae67 commit cc30810
Showing 1 changed file with 31 additions and 9 deletions.
40 changes: 31 additions & 9 deletions backend/app/agents/wazuh/services/vulnerabilities.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand All @@ -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.
Expand All @@ -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.
Expand Down

0 comments on commit cc30810

Please sign in to comment.