precommit-fixes
Taylor committed
May 29, 2024 at 09:39 UTC
6c1134359852fab1ec8a53e27cd0d15b03a26aa1
4 files changed
+18
-9
backend/app/agents/routes/agents.py
+5
-2
@@ -11,8 +11,9 @@ from sqlalchemy.ext.asyncio import AsyncSession
11
from sqlalchemy.future import select
12
13
from app.agents.dfir_iris.services.cases import collect_agent_soc_cases
14
-from app.agents.schema.agents import AgentModifyResponse, AgentWazuhUpgradeResponse
14
+from app.agents.schema.agents import AgentModifyResponse
15
from app.agents.schema.agents import AgentsResponse
16
+from app.agents.schema.agents import AgentWazuhUpgradeResponse
17
from app.agents.schema.agents import OutdatedVelociraptorAgentsResponse
18
from app.agents.schema.agents import OutdatedWazuhAgentsResponse
19
from app.agents.schema.agents import SyncedAgentsResponse
@@ -24,7 +25,8 @@ from app.agents.velociraptor.services.agents import delete_agent_velociraptor
25
from app.agents.wazuh.schema.agents import WazuhAgentScaPolicyResultsResponse
26
from app.agents.wazuh.schema.agents import WazuhAgentScaResponse
27
from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilitiesResponse
27
-from app.agents.wazuh.services.agents import delete_agent_wazuh, upgrade_wazuh_agent
28
+from app.agents.wazuh.services.agents import delete_agent_wazuh
29
+from app.agents.wazuh.services.agents import upgrade_wazuh_agent
30
from app.agents.wazuh.services.sca import collect_agent_sca
31
from app.agents.wazuh.services.sca import collect_agent_sca_policy_results
32
from app.agents.wazuh.services.vulnerabilities import collect_agent_vulnerabilities
@@ -348,6 +350,7 @@ async def mark_agent_as_not_critical(
350
detail=f"Failed to mark agent as not critical: {str(e)}",
351
)
352
353
+
354
@agents_router.post(
355
"/{agent_id}/wazuh/upgrade",
356
response_model=AgentModifyResponse,
backend/app/agents/schema/agents.py
+1
@@ -53,6 +53,7 @@ class AgentUpdateCustomerCodeResponse(BaseModel):
53
success: bool
54
message: str
55
56
+
57
class AgentWazuhUpgradeResponse(BaseModel):
58
success: bool
59
message: str
backend/app/agents/wazuh/services/agents.py
+11
-7
@@ -3,11 +3,13 @@ import asyncio
3
from fastapi import HTTPException
4
from loguru import logger
5
6
-from app.agents.schema.agents import AgentModifyResponse, AgentWazuhUpgradeResponse
6
+from app.agents.schema.agents import AgentModifyResponse
7
+from app.agents.schema.agents import AgentWazuhUpgradeResponse
8
from app.agents.wazuh.schema.agents import WazuhAgent
9
from app.agents.wazuh.schema.agents import WazuhAgentsList
10
from app.connectors.wazuh_manager.utils.universal import send_delete_request
10
-from app.connectors.wazuh_manager.utils.universal import send_get_request, send_put_request
11
+from app.connectors.wazuh_manager.utils.universal import send_get_request
12
+from app.connectors.wazuh_manager.utils.universal import send_put_request
13
14
15
async def collect_wazuh_agents() -> WazuhAgentsList:
@@ -146,6 +148,7 @@ async def delete_agent_wazuh(agent_id: str) -> AgentModifyResponse:
148
detail=f"Failed to delete agent {agent_id} from Wazuh Manager: {e}",
149
)
150
151
+
152
def handle_agent_upgrade_response(agent_upgraded: dict) -> AgentWazuhUpgradeResponse:
153
"""
154
Handle the response from the agent upgrade request.
@@ -156,24 +159,25 @@ def handle_agent_upgrade_response(agent_upgraded: dict) -> AgentWazuhUpgradeResp
159
Returns:
160
AgentWazuhUpgradeResponse: The response indicating the status of the agent upgrade.
161
"""
159
- data = agent_upgraded.get('data', {}).get('data', {})
160
- total_failed_items = data.get('total_failed_items', 0)
162
+ data = agent_upgraded.get("data", {}).get("data", {})
163
+ total_failed_items = data.get("total_failed_items", 0)
164
165
if total_failed_items == 0:
166
# Upgrade was successful
167
return AgentWazuhUpgradeResponse(
168
success=True,
166
- message=agent_upgraded.get('data', {}).get('message', 'Unknown error'),
169
+ message=agent_upgraded.get("data", {}).get("message", "Unknown error"),
170
)
171
else:
172
# Upgrade failed
170
- failed_items = data.get('failed_items', [{}])
171
- error_message = failed_items[0].get('error', {}).get('message', 'Unknown error')
173
+ failed_items = data.get("failed_items", [{}])
174
+ error_message = failed_items[0].get("error", {}).get("message", "Unknown error")
175
return AgentWazuhUpgradeResponse(
176
success=False,
177
message=error_message,
178
)
179
180
+
181
async def upgrade_wazuh_agent(agent_id: str) -> AgentWazuhUpgradeResponse:
182
"""Upgrade agent from Wazuh Manager.
183
backend/app/integrations/monitoring_alert/routes/monitoring_alert.py
+1
@@ -188,6 +188,7 @@ async def delete_monitoring_alert(
188
message="Monitoring alert deleted successfully",
189
)
190
191
+
192
@monitoring_alerts_router.delete(
193
"/purge",
194
response_model=MonitoringAlertsResponseModel,