Add endpoint to sync agent vulnerabilities by customer code (#369)
taylor_socfortress committed
Dec 12, 2024 at 11:29 UTC
96b01a12fa9700e05ac3edf82029be2e462a0133
2 files changed
+51
backend/app/agents/routes/agents.py
+27
@@ -21,6 +21,7 @@ from app.agents.schema.agents import AgentWazuhUpgradeResponse
21
from app.agents.schema.agents import OutdatedVelociraptorAgentsResponse
22
from app.agents.schema.agents import OutdatedWazuhAgentsResponse
23
from app.agents.schema.agents import SyncedAgentsResponse
24
+from app.agents.services.status import get_agents_by_customer_code
25
from app.agents.services.status import get_outdated_agents_velociraptor
26
from app.agents.services.status import get_outdated_agents_wazuh
27
from app.agents.services.sync import sync_agents_velociraptor
@@ -820,6 +821,32 @@ async def sync_vulnerabilities_route(
821
return {"success": True, "message": "Agent vulnerabilities synced successfully"}
822
823
824
+@agents_router.get(
825
+ "/sync/vulnerabilities/{customer_code}",
826
+ description="Sync agent vulnerabilities",
827
+)
828
+async def sync_vulnerabilities_customer_code_route(
829
+ customer_code: str,
830
+ session: AsyncSession = Depends(get_db),
831
+):
832
+ """
833
+ Only applies to Wazuh Manager Version 4.8.1 or higher.
834
+ 1. Loops through all agents in the database to collect their agent_name and customer code.
835
+ 2. Queries the `wazuh-states-vulnerabilities-*` index in Wazuh Indexer to get vulnerabilities based on the agent_name.
836
+ 3. Checks the `wazuh-vulnerabilities-*customer_code*` index in Wazuh Indexer to get vulnerabilities based on the
837
+ agent_name and checks to see if a vulnerability_id already exists.
838
+ 4. If the vulnerability_id does not exist, it is sent to the Graylog GELF Input.
839
+ """
840
+ logger.info("Syncing agent vulnerabilities")
841
+ agents = await get_agents_by_customer_code(customer_code, session)
842
+ for agent in agents:
843
+ if agent.customer_code is None:
844
+ logger.info(f"Skipping agent {agent.hostname} due to missing customer code")
845
+ continue
846
+ await sync_agent_vulnerabilities(agent.hostname, customer_code)
847
+ return {"success": True, "message": "Agent vulnerabilities synced successfully"}
848
+
849
+
850
# ! TODO: CURRENTLY UPDATES IN THE DB BUT NEED TO UPDATE IN WAZUH # !
851
# @agents_router.put(
852
# "/{agent_id}/update-customer-code",
backend/app/agents/services/status.py
+24
@@ -32,6 +32,30 @@ def get_agent(agent_id: str) -> List[Agents]:
32
)
33
34
35
+async def get_agents_by_customer_code(customer_code: str, session: AsyncSession) -> List[Agents]:
36
+ """
37
+ Retrieves all agents associated with a specific customer code from the database asynchronously.
38
+
39
+ Args:
40
+ customer_code (str): The customer code to filter agents by.
41
+ session (AsyncSession): The SQLAlchemy asynchronous session to use for the query.
42
+
43
+ Returns:
44
+ List[Agents]: A list of agents associated with the customer code.
45
+ """
46
+ try:
47
+ agents_result = await session.execute(select(Agents).filter(Agents.customer_code == customer_code))
48
+ agents = agents_result.scalars().all()
49
+
50
+ return agents
51
+ except Exception as e:
52
+ logger.error(f"Failed to fetch agents with customer_code {customer_code}: {e}")
53
+ raise HTTPException(
54
+ status_code=500,
55
+ detail=f"Failed to fetch agents with customer_code {customer_code}: {e}",
56
+ )
57
+
58
+
59
async def get_agent_os_by_id(agent_id: str, session: AsyncSession) -> str:
60
"""
61
Retrieves the operating system of a specific agent from the database using its ID.