@cryptotaxi247 / CoPilot / commits / 71edca0f

Vuln overview (#503)

* Add AgentVulnerabilities model and migration script * Add customer code to AgentVulnerabilities model and related database migration - Introduced customer_code field in AgentVulnerabilities model. - Updated create_from_model method to accept customer_code. - Added migration script to create customer_code column in the database. - Implemented new routes and services for vulnerability synchronization, including customer_code handling. * feat: Add vulnerability deletion endpoint and refactor vulnerability sync logic - Introduced a new endpoint to delete vulnerabilities based on agent or customer scope. - Refactored vulnerability sync functions to streamline processing and improve error handling. - Updated schemas to include response models for vulnerability deletion. - Enhanced logging for better traceability during vulnerability sync and deletion operations. - Removed unnecessary service class instantiation in favor of direct function calls for better performance. * feat: Enhance vulnerability sync with batch and bulk processing options * feat: Add vulnerability search endpoint with filtering and pagination * feat: Enhance vulnerability search to include complete agent hostname to customer code mapping * feat: Mark several vulnerability endpoints as deprecated in favor of more specific ones * feat: Integrate EPSS scoring into vulnerability search with optional inclusion * precommit fixes

taylor_socfortress committed Sep 10, 2025 at 11:37 UTC 71edca0fe4d61f8ffc0e3b3eba142e46fcebef38
11 files changed +1821
backend/alembic/versions/7b2bbee2f3e8_add_customer_code_to_vulnerabilities_.py new
+35
@@ -0,0 +1,35 @@
1 +"""Add customer code to vulnerabilities table
2 +
3 +Revision ID: 7b2bbee2f3e8
4 +Revises: dd0e62747e2c
5 +Create Date: 2025-09-09 15:17:41.770559
6 +
7 +"""
8 +from typing import Sequence
9 +from typing import Union
10 +
11 +import sqlalchemy as sa
12 +
13 +from alembic import op
14 +
15 +# revision identifiers, used by Alembic.
16 +revision: str = "7b2bbee2f3e8"
17 +down_revision: Union[str, None] = "dd0e62747e2c"
18 +branch_labels: Union[str, Sequence[str], None] = None
19 +depends_on: Union[str, Sequence[str], None] = None
20 +
21 +
22 +def upgrade() -> None:
23 + # ### commands auto generated by Alembic - please adjust! ###
24 + op.add_column("agent_vulnerabilities", sa.Column("customer_code", sa.String(length=50), nullable=True))
25 + op.create_index(op.f("ix_agent_vulnerabilities_customer_code"), "agent_vulnerabilities", ["customer_code"], unique=False)
26 + op.create_foreign_key(None, "agent_vulnerabilities", "customers", ["customer_code"], ["customer_code"])
27 + # ### end Alembic commands ###
28 +
29 +
30 +def downgrade() -> None:
31 + # ### commands auto generated by Alembic - please adjust! ###
32 + op.drop_constraint(None, "agent_vulnerabilities", type_="foreignkey")
33 + op.drop_index(op.f("ix_agent_vulnerabilities_customer_code"), table_name="agent_vulnerabilities")
34 + op.drop_column("agent_vulnerabilities", "customer_code")
35 + # ### end Alembic commands ###
backend/alembic/versions/dd0e62747e2c_add_agent_vulnerabilities_table.py new
+60
@@ -0,0 +1,60 @@
1 +"""Add agent vulnerabilities table
2 +
3 +Revision ID: dd0e62747e2c
4 +Revises: 6796a7d001ad
5 +Create Date: 2025-09-09 14:52:29.588326
6 +
7 +"""
8 +from typing import Sequence
9 +from typing import Union
10 +
11 +import sqlalchemy as sa
12 +
13 +from alembic import op
14 +
15 +# revision identifiers, used by Alembic.
16 +revision: str = "dd0e62747e2c"
17 +down_revision: Union[str, None] = "6796a7d001ad"
18 +branch_labels: Union[str, Sequence[str], None] = None
19 +depends_on: Union[str, Sequence[str], None] = None
20 +
21 +
22 +def upgrade() -> None:
23 + # ### commands auto generated by Alembic - please adjust! ###
24 + op.create_table(
25 + "agent_vulnerabilities",
26 + sa.Column("id", sa.Integer(), nullable=False),
27 + sa.Column("cve_id", sa.String(length=50), nullable=False),
28 + sa.Column("severity", sa.String(length=50), nullable=False),
29 + sa.Column("title", sa.String(length=255), nullable=False),
30 + sa.Column("references", sa.String(length=2048), nullable=True),
31 + sa.Column("status", sa.String(length=50), nullable=False),
32 + sa.Column("discovered_at", sa.DateTime(), nullable=False),
33 + sa.Column("remediated_at", sa.DateTime(), nullable=True),
34 + sa.Column("epss_score", sa.String(length=50), nullable=True),
35 + sa.Column("epss_percentile", sa.String(length=50), nullable=True),
36 + sa.Column("package_name", sa.String(length=255), nullable=True),
37 + sa.Column("agent_id", sa.String(length=256), nullable=False),
38 + sa.ForeignKeyConstraint(
39 + ["agent_id"],
40 + ["agents.agent_id"],
41 + ),
42 + sa.PrimaryKeyConstraint("id"),
43 + )
44 + op.create_index(op.f("ix_agent_vulnerabilities_agent_id"), "agent_vulnerabilities", ["agent_id"], unique=False)
45 + op.create_index(op.f("ix_agent_vulnerabilities_cve_id"), "agent_vulnerabilities", ["cve_id"], unique=False)
46 + op.create_index(op.f("ix_agent_vulnerabilities_discovered_at"), "agent_vulnerabilities", ["discovered_at"], unique=False)
47 + op.create_index(op.f("ix_agent_vulnerabilities_severity"), "agent_vulnerabilities", ["severity"], unique=False)
48 + op.create_index(op.f("ix_agent_vulnerabilities_status"), "agent_vulnerabilities", ["status"], unique=False)
49 + # ### end Alembic commands ###
50 +
51 +
52 +def downgrade() -> None:
53 + # ### commands auto generated by Alembic - please adjust! ###
54 + op.drop_index(op.f("ix_agent_vulnerabilities_status"), table_name="agent_vulnerabilities")
55 + op.drop_index(op.f("ix_agent_vulnerabilities_severity"), table_name="agent_vulnerabilities")
56 + op.drop_index(op.f("ix_agent_vulnerabilities_discovered_at"), table_name="agent_vulnerabilities")
57 + op.drop_index(op.f("ix_agent_vulnerabilities_cve_id"), table_name="agent_vulnerabilities")
58 + op.drop_index(op.f("ix_agent_vulnerabilities_agent_id"), table_name="agent_vulnerabilities")
59 + op.drop_table("agent_vulnerabilities")
60 + # ### end Alembic commands ###
backend/app/agents/vulnerabilities/__init__.py new
+1
@@ -0,0 +1 @@
1 +# Vulnerabilities package for agent vulnerability management
backend/app/agents/vulnerabilities/routes/__init__.py new
+1
@@ -0,0 +1 @@
1 +# Routes package
backend/app/agents/vulnerabilities/routes/vulnerabilities.py new
+435
@@ -0,0 +1,435 @@
1 +from typing import List
2 +from typing import Optional
3 +
4 +from fastapi import APIRouter
5 +from fastapi import BackgroundTasks
6 +from fastapi import Depends
7 +from fastapi import HTTPException
8 +from fastapi import Query
9 +from fastapi import Security
10 +from loguru import logger
11 +from sqlalchemy.ext.asyncio import AsyncSession
12 +
13 +from app.agents.vulnerabilities.schema.vulnerabilities import (
14 + AgentVulnerabilitiesResponse,
15 +)
16 +from app.agents.vulnerabilities.schema.vulnerabilities import (
17 + VulnerabilityDeleteResponse,
18 +)
19 +from app.agents.vulnerabilities.schema.vulnerabilities import (
20 + VulnerabilitySearchResponse,
21 +)
22 +from app.agents.vulnerabilities.schema.vulnerabilities import VulnerabilityStatsResponse
23 +from app.agents.vulnerabilities.schema.vulnerabilities import VulnerabilitySyncRequest
24 +from app.agents.vulnerabilities.schema.vulnerabilities import VulnerabilitySyncResponse
25 +from app.agents.vulnerabilities.services.vulnerabilities import delete_vulnerabilities
26 +from app.agents.vulnerabilities.services.vulnerabilities import (
27 + get_vulnerabilities_by_agent,
28 +)
29 +from app.agents.vulnerabilities.services.vulnerabilities import (
30 + get_vulnerability_statistics,
31 +)
32 +from app.agents.vulnerabilities.services.vulnerabilities import (
33 + search_vulnerabilities_from_indexer,
34 +)
35 +from app.agents.vulnerabilities.services.vulnerabilities import sync_all_vulnerabilities
36 +from app.agents.vulnerabilities.services.vulnerabilities import (
37 + sync_vulnerabilities_for_agent,
38 +)
39 +from app.auth.routes.auth import AuthHandler
40 +from app.db.db_session import get_db
41 +from app.db.db_session import get_db_session
42 +
43 +# Create router for vulnerability endpoints
44 +vulnerabilities_router = APIRouter()
45 +
46 +
47 +@vulnerabilities_router.post(
48 + "/sync",
49 + response_model=VulnerabilitySyncResponse,
50 + description="Sync vulnerabilities from Wazuh Indexer indices to database for all agents with performance options",
51 + dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst", "scheduler"))],
52 + deprecated=True, # Marking this endpoint as deprecated in favor of more specific ones
53 +)
54 +async def sync_vulnerabilities(
55 + sync_request: Optional[VulnerabilitySyncRequest] = None,
56 + batch_size: int = Query(100, description="Batch size for processing (1-1000)", ge=1, le=1000),
57 + use_bulk_mode: bool = Query(False, description="Use ultra-fast bulk mode for large datasets"),
58 + db: AsyncSession = Depends(get_db),
59 +) -> VulnerabilitySyncResponse:
60 + """
61 + Sync vulnerabilities from Wazuh Indexer indices to the database for all agents.
62 +
63 + This endpoint fetches vulnerability data from the 'wazuh-states-vulnerabilities-*'
64 + indices for all agents in the database, processes them, and stores them in the
65 + agent_vulnerabilities table.
66 +
67 + **Performance Modes:**
68 + - **Batch Mode** (default): Processes vulnerabilities in configurable batches with individual error handling
69 + - **Bulk Mode**: Ultra-fast processing using bulk database operations for large datasets
70 +
71 + The endpoint automatically discovers all agents from the database and syncs
72 + vulnerabilities for each one using their hostname and customer_code.
73 +
74 + Args:
75 + sync_request: Optional request parameters for vulnerability sync
76 + batch_size: Number of vulnerabilities to process per batch (1-1000, default: 100)
77 + use_bulk_mode: Enable ultra-fast bulk processing mode for large datasets
78 + db: Database session
79 +
80 + Returns:
81 + VulnerabilitySyncResponse: Status of the sync operation
82 + """
83 + logger.info(f"Starting vulnerability sync for all agents from database (batch_size={batch_size}, bulk_mode={use_bulk_mode})")
84 +
85 + try:
86 + # Use the standalone function directly with performance options
87 + result = await sync_all_vulnerabilities(db_session=db, customer_code=None, batch_size=batch_size, use_bulk_mode=use_bulk_mode)
88 + return result
89 +
90 + except Exception as e:
91 + logger.error(f"Error in vulnerability sync endpoint: {e}")
92 + raise HTTPException(status_code=500, detail=f"Vulnerability sync failed: {e}")
93 +
94 +
95 +@vulnerabilities_router.post(
96 + "/sync/background",
97 + response_model=dict,
98 + description="Start background vulnerability sync for all agents",
99 + dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],
100 + deprecated=True, # Marking this endpoint as deprecated in favor of more specific ones
101 +)
102 +async def sync_vulnerabilities_background(
103 + background_tasks: BackgroundTasks,
104 + sync_request: Optional[VulnerabilitySyncRequest] = None,
105 + db: AsyncSession = Depends(get_db),
106 +):
107 + """
108 + Start vulnerability sync as a background task for all agents in the database.
109 +
110 + This endpoint automatically discovers all agents from the database and starts
111 + a background task to sync vulnerabilities for each one. This is useful for
112 + large sync operations that might take a long time.
113 +
114 + Args:
115 + background_tasks: FastAPI background tasks
116 + sync_request: Optional request parameters
117 + db: Database session
118 + """
119 + logger.info("Starting background vulnerability sync for all agents from database")
120 +
121 + async def background_sync():
122 + try:
123 + # Create a new database session for the background task
124 + async with get_db_session() as bg_db:
125 + # Use the standalone function directly with default performance settings
126 + await sync_all_vulnerabilities(db_session=bg_db, customer_code=None, batch_size=100, use_bulk_mode=False)
127 + except Exception as e:
128 + logger.error(f"Background vulnerability sync failed: {e}")
129 +
130 + background_tasks.add_task(background_sync)
131 +
132 + return {"success": True, "message": "Vulnerability sync started in background for all agents"}
133 +
134 +
135 +@vulnerabilities_router.get(
136 + "/agent/{agent_id}",
137 + response_model=AgentVulnerabilitiesResponse,
138 + description="Get vulnerabilities for a specific agent",
139 + dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],
140 + deprecated=True, # Marking this endpoint as deprecated in favor of more specific ones
141 +)
142 +async def get_agent_vulnerabilities(
143 + agent_id: str,
144 + severity: Optional[List[str]] = Query(None, description="Filter by severity levels"),
145 + db: AsyncSession = Depends(get_db),
146 +) -> AgentVulnerabilitiesResponse:
147 + """
148 + Retrieve vulnerabilities for a specific agent from the database.
149 +
150 + Args:
151 + agent_id: ID of the agent to get vulnerabilities for
152 + severity: Optional list of severity levels to filter by (Critical, High, Medium, Low)
153 + db: Database session
154 +
155 + Returns:
156 + AgentVulnerabilitiesResponse: List of vulnerabilities for the agent
157 + """
158 + logger.info(f"Getting vulnerabilities for agent: {agent_id}")
159 +
160 + try:
161 + # Use the standalone function directly
162 + return await get_vulnerabilities_by_agent(db_session=db, agent_id=agent_id, severity_filter=severity)
163 +
164 + except Exception as e:
165 + logger.error(f"Error getting vulnerabilities for agent {agent_id}: {e}")
166 + raise HTTPException(status_code=500, detail=f"Failed to get vulnerabilities for agent {agent_id}: {e}")
167 +
168 +
169 +@vulnerabilities_router.get(
170 + "/stats",
171 + response_model=VulnerabilityStatsResponse,
172 + description="Get vulnerability statistics",
173 + dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],
174 + deprecated=True, # Marking this endpoint as deprecated in favor of more specific ones
175 +)
176 +async def get_vulnerability_stats(
177 + customer_code: Optional[str] = Query(None, description="Filter by customer code"),
178 + db: AsyncSession = Depends(get_db),
179 +) -> VulnerabilityStatsResponse:
180 + """
181 + Get vulnerability statistics across all agents or for a specific customer.
182 +
183 + Args:
184 + customer_code: Optional customer code to filter statistics by
185 + db: Database session
186 +
187 + Returns:
188 + VulnerabilityStatsResponse: Vulnerability statistics
189 + """
190 + logger.info(f"Getting vulnerability statistics for customer: {customer_code}")
191 +
192 + try:
193 + # Use the standalone function directly
194 + return await get_vulnerability_statistics(db_session=db, customer_code=customer_code)
195 +
196 + except Exception as e:
197 + logger.error(f"Error getting vulnerability statistics: {e}")
198 + raise HTTPException(status_code=500, detail=f"Failed to get vulnerability statistics: {e}")
199 +
200 +
201 +@vulnerabilities_router.post(
202 + "/sync/customer/{customer_code}",
203 + response_model=VulnerabilitySyncResponse,
204 + description="Sync vulnerabilities for all agents of a specific customer",
205 + dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],
206 + deprecated=True, # Marking this endpoint as deprecated in favor of more specific ones
207 +)
208 +async def sync_customer_vulnerabilities(
209 + customer_code: str,
210 + background_tasks: BackgroundTasks,
211 + force_refresh: bool = Query(False, description="Force refresh of existing vulnerabilities"),
212 + batch_size: int = Query(100, description="Batch size for processing vulnerabilities"),
213 + use_bulk_mode: bool = Query(False, description="Use ultra-fast bulk operations for large datasets"),
214 + db: AsyncSession = Depends(get_db),
215 +) -> VulnerabilitySyncResponse:
216 + """
217 + Sync vulnerabilities for all agents belonging to a specific customer with performance options.
218 +
219 + Args:
220 + customer_code: Customer code to sync vulnerabilities for
221 + background_tasks: FastAPI background tasks
222 + force_refresh: Whether to force refresh of existing vulnerabilities
223 + batch_size: Number of vulnerabilities to process in each batch
224 + use_bulk_mode: Use ultra-fast bulk operations for large datasets
225 + db: Database session
226 +
227 + Returns:
228 + VulnerabilitySyncResponse: Status of the sync operation
229 + """
230 + logger.info(f"Starting vulnerability sync for customer: {customer_code} (batch_size={batch_size}, bulk_mode={use_bulk_mode})")
231 +
232 + try:
233 + # Use the standalone function directly with performance options
234 + result = await sync_all_vulnerabilities(
235 + db_session=db,
236 + customer_code=customer_code,
237 + batch_size=batch_size,
238 + use_bulk_mode=use_bulk_mode,
239 + )
240 + return result
241 +
242 + except Exception as e:
243 + logger.error(f"Error syncing vulnerabilities for customer {customer_code}: {e}")
244 + raise HTTPException(status_code=500, detail=f"Failed to sync vulnerabilities for customer {customer_code}: {e}")
245 +
246 +
247 +@vulnerabilities_router.post(
248 + "/sync/agent/{agent_name}",
249 + response_model=VulnerabilitySyncResponse,
250 + description="Sync vulnerabilities for a specific agent with performance options",
251 + dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],
252 + deprecated=True, # Marking this endpoint as deprecated in favor of more specific ones
253 +)
254 +async def sync_agent_vulnerabilities(
255 + agent_name: str,
256 + customer_code: Optional[str] = Query(None, description="Override customer code"),
257 + batch_size: int = Query(100, description="Batch size for processing (1-1000)", ge=1, le=1000),
258 + use_bulk_mode: bool = Query(False, description="Use ultra-fast bulk mode for large datasets"),
259 + db: AsyncSession = Depends(get_db),
260 +) -> VulnerabilitySyncResponse:
261 + """
262 + Sync vulnerabilities for a specific agent with performance optimization options.
263 +
264 + **Performance Modes:**
265 + - **Batch Mode** (default): Processes vulnerabilities in configurable batches with individual error handling
266 + - **Bulk Mode**: Ultra-fast processing using bulk database operations for large datasets
267 +
268 + Args:
269 + agent_name: Name/hostname of the agent to sync vulnerabilities for
270 + customer_code: Optional customer code override
271 + batch_size: Number of vulnerabilities to process per batch (1-1000, default: 100)
272 + use_bulk_mode: Enable ultra-fast bulk processing mode for large datasets
273 + db: Database session
274 +
275 + Returns:
276 + VulnerabilitySyncResponse: Status of the sync operation
277 + """
278 + logger.info(f"Starting vulnerability sync for agent: {agent_name} (batch_size={batch_size}, bulk_mode={use_bulk_mode})")
279 +
280 + try:
281 + # Use the standalone function directly with performance options
282 + result = await sync_vulnerabilities_for_agent(
283 + db_session=db,
284 + agent_name=agent_name,
285 + customer_code=customer_code,
286 + batch_size=batch_size,
287 + use_bulk_mode=use_bulk_mode,
288 + )
289 + return result
290 +
291 + except Exception as e:
292 + logger.error(f"Error syncing vulnerabilities for agent {agent_name}: {e}")
293 + raise HTTPException(status_code=500, detail=f"Failed to sync vulnerabilities for agent {agent_name}: {e}")
294 +
295 +
296 +@vulnerabilities_router.delete(
297 + "/delete",
298 + response_model=VulnerabilityDeleteResponse,
299 + description="Delete vulnerabilities based on scope",
300 + dependencies=[Security(AuthHandler().require_any_scope("admin"))],
301 +)
302 +async def delete_vulnerabilities_endpoint(
303 + agent_name: Optional[str] = Query(None, description="Delete vulnerabilities for specific agent"),
304 + customer_code: Optional[str] = Query(None, description="Delete vulnerabilities for specific customer"),
305 + confirm_delete_all: bool = Query(False, description="Required confirmation to delete ALL vulnerabilities"),
306 + db: AsyncSession = Depends(get_db),
307 +) -> VulnerabilityDeleteResponse:
308 + """
309 + Delete vulnerabilities based on scope:
310 +
311 + - If neither agent_name nor customer_code provided: Delete ALL vulnerabilities (requires confirm_delete_all=true)
312 + - If agent_name provided: Delete vulnerabilities for that specific agent
313 + - If customer_code provided: Delete vulnerabilities for all agents of that customer
314 +
315 + **WARNING**: Deleting all vulnerabilities is irreversible. Use with caution.
316 +
317 + Args:
318 + agent_name: Optional agent name to delete vulnerabilities for
319 + customer_code: Optional customer code to delete vulnerabilities for
320 + confirm_delete_all: Required confirmation when deleting ALL vulnerabilities
321 + db: Database session
322 +
323 + Returns:
324 + VulnerabilityDeleteResponse: Status of the delete operation
325 + """
326 +
327 + # Safety check for deleting all vulnerabilities
328 + if not agent_name and not customer_code:
329 + if not confirm_delete_all:
330 + raise HTTPException(status_code=400, detail="To delete ALL vulnerabilities, you must set confirm_delete_all=true")
331 + logger.warning("Request to delete ALL vulnerabilities received with confirmation")
332 +
333 + # Validate that both agent_name and customer_code are not provided
334 + if agent_name and customer_code:
335 + raise HTTPException(status_code=400, detail="Cannot specify both agent_name and customer_code. Choose one scope.")
336 +
337 + try:
338 + result = await delete_vulnerabilities(db_session=db, agent_name=agent_name, customer_code=customer_code)
339 + return result
340 +
341 + except Exception as e:
342 + logger.error(f"Error in delete vulnerabilities endpoint: {e}")
343 + raise HTTPException(status_code=500, detail=f"Failed to delete vulnerabilities: {e}")
344 +
345 +
346 +@vulnerabilities_router.get(
347 + "/search",
348 + response_model=VulnerabilitySearchResponse,
349 + description="Search vulnerabilities directly from Wazuh indexer with filtering and pagination",
350 + dependencies=[Security(AuthHandler().require_any_scope("admin", "analyst"))],
351 +)
352 +async def search_vulnerabilities(
353 + customer_code: Optional[str] = Query(None, description="Filter by customer code"),
354 + agent_name: Optional[str] = Query(None, description="Filter by agent hostname"),
355 + severity: Optional[str] = Query(None, description="Filter by severity (Critical, High, Medium, Low)"),
356 + cve_id: Optional[str] = Query(None, description="Filter by specific CVE ID"),
357 + package_name: Optional[str] = Query(None, description="Filter by package name"),
358 + page: int = Query(1, description="Page number for pagination", ge=1),
359 + page_size: int = Query(50, description="Number of vulnerabilities per page", ge=1, le=1000),
360 + include_epss: bool = Query(True, description="Include EPSS scores (may impact performance)"),
361 + db: AsyncSession = Depends(get_db),
362 +) -> VulnerabilitySearchResponse:
363 + """
364 + Search vulnerabilities directly from Wazuh indexer without storing them in database.
365 +
366 + This endpoint provides fast, real-time vulnerability data with advanced filtering
367 + and pagination capabilities. Perfect for exploring vulnerability data without
368 + the overhead of database synchronization.
369 +
370 + **Features:**
371 + - Real-time data directly from Wazuh indexer
372 + - Advanced filtering by customer, agent, severity, CVE, or package
373 + - Efficient pagination for large result sets
374 + - No database storage required
375 + - Optional EPSS scoring integration
376 +
377 + **Performance:**
378 + - Handles large datasets efficiently with pagination
379 + - Optimized Elasticsearch queries for fast response times
380 + - Automatic sorting by detection date and severity
381 + - EPSS scoring can be disabled for faster response times
382 +
383 + **Filtering Options:**
384 + - **customer_code**: Filter by specific customer
385 + - **agent_name**: Filter by specific agent hostname
386 + - **severity**: Filter by vulnerability severity (Critical, High, Medium, Low)
387 + - **cve_id**: Search for specific CVE identifier
388 + - **package_name**: Filter by package name (supports partial matching)
389 +
390 + **EPSS Integration:**
391 + - **include_epss**: Include EPSS scores and percentiles for vulnerabilities
392 + - Provides risk assessment data from FIRST.org
393 + - May impact response time due to external API calls
394 +
395 + **Pagination:**
396 + - **page**: Page number (starts at 1)
397 + - **page_size**: Results per page (1-1000, default: 50)
398 +
399 + Args:
400 + customer_code: Optional customer code filter
401 + agent_name: Optional agent hostname filter
402 + severity: Optional severity filter
403 + cve_id: Optional CVE ID filter
404 + package_name: Optional package name filter (partial matching)
405 + page: Page number for pagination
406 + page_size: Number of results per page
407 + db: Database session
408 +
409 + Returns:
410 + VulnerabilitySearchResponse: Paginated vulnerability search results
411 + """
412 + logger.info(
413 + f"Searching vulnerabilities from indexer with filters: "
414 + f"customer_code={customer_code}, agent_name={agent_name}, "
415 + f"severity={severity}, cve_id={cve_id}, package_name={package_name}, "
416 + f"page={page}, page_size={page_size}, include_epss={include_epss}",
417 + )
418 +
419 + try:
420 + result = await search_vulnerabilities_from_indexer(
421 + db_session=db,
422 + customer_code=customer_code,
423 + agent_name=agent_name,
424 + severity=severity,
425 + cve_id=cve_id,
426 + package_name=package_name,
427 + page=page,
428 + page_size=page_size,
429 + include_epss=include_epss,
430 + )
431 + return result
432 +
433 + except Exception as e:
434 + logger.error(f"Error in vulnerability search endpoint: {e}")
435 + raise HTTPException(status_code=500, detail=f"Failed to search vulnerabilities: {e}")
backend/app/agents/vulnerabilities/schema/__init__.py new
+1
@@ -0,0 +1 @@
1 +# Schema package
backend/app/agents/vulnerabilities/schema/vulnerabilities.py new
+136
@@ -0,0 +1,136 @@
1 +from datetime import datetime
2 +from typing import List
3 +from typing import Optional
4 +
5 +from pydantic import BaseModel
6 +from pydantic import Field
7 +
8 +
9 +class WazuhVulnerabilityData(BaseModel):
10 + """Schema for processing raw Wazuh vulnerability data from Elasticsearch"""
11 +
12 + cve_id: str = Field(..., alias="id")
13 + severity: str
14 + title: str = Field(..., alias="description")
15 + references: Optional[str] = Field(None, alias="reference")
16 + detected_at: datetime
17 + published_at: Optional[datetime] = None
18 + base_score: Optional[float] = Field(None, alias="base")
19 + package_name: Optional[str] = None
20 + package_version: Optional[str] = None
21 + package_architecture: Optional[str] = None
22 +
23 + class Config:
24 + allow_population_by_field_name = True
25 +
26 +
27 +class AgentVulnerabilityOut(BaseModel):
28 + """Output schema for agent vulnerabilities"""
29 +
30 + id: int
31 + cve_id: str
32 + severity: str
33 + title: str
34 + references: Optional[str] = None
35 + status: str
36 + discovered_at: datetime
37 + remediated_at: Optional[datetime] = None
38 + epss_score: Optional[str] = None
39 + epss_percentile: Optional[str] = None
40 + package_name: Optional[str] = None
41 + agent_id: str
42 + customer_code: Optional[str] = None
43 +
44 +
45 +class AgentVulnerabilitiesResponse(BaseModel):
46 + """Response schema for agent vulnerabilities"""
47 +
48 + vulnerabilities: List[AgentVulnerabilityOut]
49 + success: bool
50 + message: str
51 + total_count: int
52 +
53 +
54 +class VulnerabilitySyncRequest(BaseModel):
55 + """Request schema for syncing vulnerabilities - all fields optional"""
56 +
57 + customer_code: Optional[str] = Field(None, description="Optional customer code filter")
58 + agent_name: Optional[str] = Field(None, description="Optional specific agent name")
59 + force_refresh: bool = Field(False, description="Force refresh of existing vulnerabilities")
60 +
61 +
62 +class VulnerabilitySyncResponse(BaseModel):
63 + """Response schema for vulnerability sync operations"""
64 +
65 + success: bool
66 + message: str
67 + synced_count: int
68 + errors: List[str] = []
69 +
70 +
71 +class VulnerabilityStatsResponse(BaseModel):
72 + """Response schema for vulnerability statistics"""
73 +
74 + total_vulnerabilities: int
75 + critical_count: int
76 + high_count: int
77 + medium_count: int
78 + low_count: int
79 + by_customer: dict = {}
80 + success: bool
81 + message: str
82 +
83 +
84 +class VulnerabilityDeleteResponse(BaseModel):
85 + """Response schema for vulnerability delete operations"""
86 +
87 + success: bool
88 + message: str
89 + deleted_count: int
90 + errors: List[str] = []
91 +
92 +
93 +class VulnerabilitySearchRequest(BaseModel):
94 + """Request schema for searching vulnerabilities from Wazuh indexer"""
95 +
96 + customer_code: Optional[str] = Field(None, description="Filter by customer code")
97 + agent_name: Optional[str] = Field(None, description="Filter by agent hostname")
98 + severity: Optional[str] = Field(None, description="Filter by severity (Critical, High, Medium, Low)")
99 + page: int = Field(1, description="Page number for pagination", ge=1)
100 + page_size: int = Field(50, description="Number of vulnerabilities per page", ge=1, le=1000)
101 + cve_id: Optional[str] = Field(None, description="Filter by specific CVE ID")
102 + package_name: Optional[str] = Field(None, description="Filter by package name")
103 +
104 +
105 +class VulnerabilitySearchItem(BaseModel):
106 + """Individual vulnerability item from search results"""
107 +
108 + cve_id: str
109 + severity: str
110 + title: str
111 + agent_name: str
112 + customer_code: Optional[str] = None
113 + references: Optional[str] = None
114 + detected_at: datetime
115 + published_at: Optional[datetime] = None
116 + base_score: Optional[float] = None
117 + package_name: Optional[str] = None
118 + package_version: Optional[str] = None
119 + package_architecture: Optional[str] = None
120 + epss_score: Optional[str] = None
121 + epss_percentile: Optional[str] = None
122 +
123 +
124 +class VulnerabilitySearchResponse(BaseModel):
125 + """Response schema for vulnerability search results with pagination"""
126 +
127 + vulnerabilities: List[VulnerabilitySearchItem]
128 + total_count: int
129 + page: int
130 + page_size: int
131 + total_pages: int
132 + has_next: bool
133 + has_previous: bool
134 + success: bool
135 + message: str
136 + filters_applied: dict = {}
backend/app/agents/vulnerabilities/services/__init__.py new
+1
@@ -0,0 +1 @@
1 +# Services package
backend/app/agents/vulnerabilities/services/vulnerabilities.py new
+1085
@@ -0,0 +1,1085 @@
1 +from datetime import datetime
2 +from typing import Any
3 +from typing import Dict
4 +from typing import List
5 +from typing import Optional
6 +
7 +from fastapi import HTTPException
8 +from loguru import logger
9 +from sqlalchemy import select
10 +from sqlalchemy.ext.asyncio import AsyncSession
11 +
12 +from app.agents.vulnerabilities.schema.vulnerabilities import (
13 + AgentVulnerabilitiesResponse,
14 +)
15 +from app.agents.vulnerabilities.schema.vulnerabilities import AgentVulnerabilityOut
16 +from app.agents.vulnerabilities.schema.vulnerabilities import VulnerabilitySearchItem
17 +from app.agents.vulnerabilities.schema.vulnerabilities import (
18 + VulnerabilitySearchResponse,
19 +)
20 +from app.agents.vulnerabilities.schema.vulnerabilities import VulnerabilityStatsResponse
21 +from app.agents.vulnerabilities.schema.vulnerabilities import VulnerabilitySyncResponse
22 +from app.agents.vulnerabilities.schema.vulnerabilities import WazuhVulnerabilityData
23 +from app.connectors.wazuh_indexer.utils.universal import collect_indices
24 +from app.connectors.wazuh_indexer.utils.universal import (
25 + create_wazuh_indexer_client_async,
26 +)
27 +from app.db.universal_models import Agents
28 +from app.db.universal_models import AgentVulnerabilities
29 +from app.threat_intel.schema.epss import EpssThreatIntelRequest
30 +from app.threat_intel.services.epss import collect_epss_score
31 +
32 +
33 +async def get_epss_score_for_cve(cve_id: str) -> tuple[Optional[str], Optional[str]]:
34 + """
35 + Get EPSS score and percentile for a CVE ID
36 +
37 + Args:
38 + cve_id: CVE identifier to get EPSS score for
39 +
40 + Returns:
41 + Tuple of (epss_score, epss_percentile) or (None, None) if not found
42 + """
43 + try:
44 + epss_request = EpssThreatIntelRequest(cve=cve_id)
45 + epss_response = await collect_epss_score(epss_request)
46 +
47 + if epss_response.success and epss_response.data:
48 + # Get the first (and usually only) result
49 + epss_data = epss_response.data[0]
50 + return epss_data.epss, epss_data.percentile
51 + else:
52 + logger.debug(f"No EPSS data found for CVE: {cve_id}")
53 + return None, None
54 +
55 + except Exception as e:
56 + logger.warning(f"Error fetching EPSS score for CVE {cve_id}: {e}")
57 + return None, None
58 +
59 +
60 +def process_wazuh_document(document: Dict[str, Any]) -> WazuhVulnerabilityData:
61 + """
62 + Process a single Wazuh vulnerability document from Indexer
63 +
64 + Args:
65 + document: Raw document from Wazuh Indexer index
66 +
67 + Returns:
68 + WazuhVulnerabilityData: Processed vulnerability data
69 + """
70 + logger.info(f"Processing vulnerability document ID: {document.get('_id', 'unknown')}")
71 + try:
72 + source = document.get("_source", {})
73 + vuln_data = source.get("vulnerability", {})
74 + package_data = source.get("package", {})
75 + score_data = vuln_data.get("score", {})
76 +
77 + # Parse detected_at timestamp
78 + detected_at_str = vuln_data.get("detected_at")
79 + detected_at = datetime.fromisoformat(detected_at_str.replace("Z", "+00:00")) if detected_at_str else datetime.utcnow()
80 +
81 + # Parse published_at timestamp if available
82 + published_at_str = vuln_data.get("published_at")
83 + published_at = None
84 + if published_at_str:
85 + try:
86 + published_at = datetime.fromisoformat(published_at_str.replace("Z", "+00:00"))
87 + except ValueError:
88 + logger.warning(f"Could not parse published_at: {published_at_str}")
89 +
90 + # Parse and limit references to first 5 items if comma-separated
91 + references_raw = vuln_data.get("reference")
92 + references = None
93 + if references_raw:
94 + if isinstance(references_raw, str) and "," in references_raw:
95 + # Split by comma, take first 5 items, and rejoin
96 + reference_list = [ref.strip() for ref in references_raw.split(",")]
97 + references = ", ".join(reference_list[:5])
98 + else:
99 + references = str(references_raw)
100 +
101 + # Also ensure the references field doesn't exceed database column limit (2048 chars)
102 + if len(references) > 2048:
103 + references = references[:2045] + "..."
104 +
105 + return WazuhVulnerabilityData(
106 + cve_id=vuln_data.get("id", "UNKNOWN_CVE"),
107 + severity=vuln_data.get("severity", "UNKNOWN"),
108 + title=package_data.get("name", "Unknown Package"),
109 + references=references,
110 + detected_at=detected_at,
111 + published_at=published_at,
112 + base_score=score_data.get("base"),
113 + package_name=package_data.get("name"),
114 + package_version=package_data.get("version"),
115 + package_architecture=package_data.get("architecture"),
116 + )
117 + except Exception as e:
118 + logger.error(f"Error processing vulnerability document: {e}")
119 + logger.error(f"Document: {document}")
120 + raise
121 +
122 +
123 +async def get_vulnerabilities_indices() -> List[str]:
124 + """Get all vulnerability indices from Wazuh Indexer"""
125 + try:
126 + indices = await collect_indices(all_indices=True)
127 + vuln_indices = [index for index in indices.indices_list if index.startswith("wazuh-states-vulnerabilities")]
128 + logger.info(f"Found {len(vuln_indices)} vulnerability indices")
129 + return vuln_indices
130 + except Exception as e:
131 + logger.error(f"Error collecting vulnerability indices: {e}")
132 + raise HTTPException(status_code=500, detail=f"Failed to collect vulnerability indices: {e}")
133 +
134 +
135 +async def fetch_vulnerabilities_from_indexer(
136 + agent_name: Optional[str] = None,
137 + customer_code: Optional[str] = None,
138 + severity_filter: Optional[List[str]] = None,
139 +) -> List[Dict[str, Any]]:
140 + """
141 + Fetch vulnerabilities from Wazuh Indexer indices
142 +
143 + Args:
144 + agent_name: Optional agent name filter
145 + customer_code: Optional customer code filter (used for index filtering)
146 + severity_filter: Optional list of severities to filter by
147 +
148 + Returns:
149 + List of vulnerability documents
150 + """
151 + es_client = None
152 + try:
153 + es_client = await create_wazuh_indexer_client_async("Wazuh-Indexer")
154 + indices = await get_vulnerabilities_indices()
155 +
156 + if not indices:
157 + logger.warning("No vulnerability indices found")
158 + return []
159 +
160 + vulnerabilities = []
161 +
162 + # Build query
163 + query = {"query": {"bool": {"must": []}}}
164 +
165 + if agent_name:
166 + query["query"]["bool"]["must"].append({"match": {"agent.name": agent_name}})
167 +
168 + if severity_filter:
169 + query["query"]["bool"]["must"].append({"terms": {"vulnerability.severity": severity_filter}})
170 +
171 + # If no filters, match all
172 + if not query["query"]["bool"]["must"]:
173 + query = {"query": {"match_all": {}}}
174 +
175 + # Search across all vulnerability indices
176 + for index in indices:
177 + try:
178 + # Use scroll for large result sets
179 + page = await es_client.search(index=index, body=query, scroll="2m", size=1000)
180 + scroll_id = page["_scroll_id"]
181 + scroll_size = len(page["hits"]["hits"])
182 +
183 + vulnerabilities.extend(page["hits"]["hits"])
184 +
185 + # Continue scrolling through results
186 + while scroll_size > 0:
187 + page = await es_client.scroll(scroll_id=scroll_id, scroll="2m")
188 + scroll_id = page["_scroll_id"]
189 + scroll_size = len(page["hits"]["hits"])
190 + vulnerabilities.extend(page["hits"]["hits"])
191 +
192 + # Clear the scroll context when done with this index
193 + try:
194 + await es_client.clear_scroll(scroll_id=scroll_id)
195 + except Exception as clear_error:
196 + logger.warning(f"Could not clear scroll context: {clear_error}")
197 +
198 + except Exception as index_error:
199 + logger.error(f"Error querying index {index}: {index_error}")
200 + continue
201 +
202 + logger.info(f"Fetched {len(vulnerabilities)} vulnerabilities from Indexer")
203 + return vulnerabilities
204 +
205 + except Exception as e:
206 + logger.error(f"Error fetching vulnerabilities from Indexer: {e}")
207 + raise HTTPException(status_code=500, detail=f"Failed to fetch vulnerabilities: {e}")
208 + finally:
209 + # Ensure the Elasticsearch client session is properly closed
210 + if es_client:
211 + try:
212 + await es_client.close()
213 + except Exception as close_error:
214 + logger.warning(f"Error closing Elasticsearch client: {close_error}")
215 +
216 +
217 +async def get_agent_by_name(db_session: AsyncSession, agent_name: str) -> Optional[Agents]:
218 + """Get agent from database by hostname/name"""
219 + try:
220 + result = await db_session.execute(select(Agents).filter(Agents.hostname == agent_name))
221 + return result.scalars().first()
222 + except Exception as e:
223 + logger.error(f"Error fetching agent {agent_name}: {e}")
224 + return None
225 +
226 +
227 +async def _sync_vulnerabilities_bulk_mode(
228 + db_session: AsyncSession,
229 + agent_id: str,
230 + agent_name: str,
231 + customer_code: str,
232 + vulnerability_docs: List[Dict[str, Any]],
233 +) -> "VulnerabilitySyncResponse":
234 + """
235 + Ultra-fast bulk mode for processing large numbers of vulnerabilities.
236 + Uses SQLAlchemy bulk operations for maximum performance.
237 + """
238 + from app.agents.vulnerabilities.schema.vulnerabilities import (
239 + VulnerabilitySyncResponse,
240 + )
241 +
242 + try:
243 + logger.info(f"BULK MODE: Processing {len(vulnerability_docs)} vulnerabilities for agent {agent_name}")
244 +
245 + # Process all documents first
246 + processed_vulns = []
247 + errors = []
248 +
249 + for doc in vulnerability_docs:
250 + try:
251 + vuln_data = process_wazuh_document(doc)
252 + processed_vulns.append(vuln_data)
253 + except Exception as doc_error:
254 + error_msg = f"Error processing vulnerability {doc.get('_id', 'unknown')}: {doc_error}"
255 + logger.error(error_msg)
256 + errors.append(error_msg)
257 +
258 + if not processed_vulns:
259 + return VulnerabilitySyncResponse(
260 + success=True,
261 + message=f"No valid vulnerabilities to process for agent {agent_name}",
262 + synced_count=0,
263 + errors=errors,
264 + )
265 +
266 + # Get existing vulnerabilities for comparison
267 + existing_vulns_result = await db_session.execute(select(AgentVulnerabilities).filter(AgentVulnerabilities.agent_id == agent_id))
268 + existing_vulns = existing_vulns_result.scalars().all()
269 +
270 + # Create lookup for existing vulnerabilities
271 + existing_lookup = {}
272 + for vuln in existing_vulns:
273 + key = f"{vuln.cve_id}_{vuln.package_name or 'None'}"
274 + existing_lookup[key] = vuln
275 +
276 + # Prepare bulk operations
277 + new_vulnerabilities = []
278 + update_data = []
279 +
280 + for vuln_data in processed_vulns:
281 + key = f"{vuln_data.cve_id}_{vuln_data.package_name or 'None'}"
282 +
283 + if key in existing_lookup:
284 + # Prepare for bulk update
285 + existing_vuln = existing_lookup[key]
286 + update_data.append(
287 + {
288 + "id": existing_vuln.id,
289 + "severity": vuln_data.severity,
290 + "title": vuln_data.title,
291 + "references": vuln_data.references,
292 + "discovered_at": vuln_data.detected_at,
293 + "epss_score": str(vuln_data.base_score)
294 + if hasattr(vuln_data, "base_score") and vuln_data.base_score
295 + else existing_vuln.epss_score,
296 + "package_name": vuln_data.package_name,
297 + },
298 + )
299 + else:
300 + # Prepare for bulk insert
301 + new_vuln = AgentVulnerabilities.create_from_model(
302 + vulnerability_data=vuln_data,
303 + agent_id=agent_id,
304 + customer_code=customer_code,
305 + )
306 + new_vulnerabilities.append(new_vuln)
307 +
308 + # Execute bulk operations
309 + inserted_count = 0
310 + updated_count = 0
311 +
312 + if new_vulnerabilities:
313 + db_session.add_all(new_vulnerabilities)
314 + inserted_count = len(new_vulnerabilities)
315 + logger.info(f"BULK MODE: Prepared {inserted_count} new vulnerabilities for insertion")
316 +
317 + if update_data:
318 + # Use bulk update for existing vulnerabilities
319 + from sqlalchemy import update
320 +
321 + for data in update_data:
322 + stmt = (
323 + update(AgentVulnerabilities)
324 + .where(AgentVulnerabilities.id == data["id"])
325 + .values(
326 + {
327 + "severity": data["severity"],
328 + "title": data["title"],
329 + "references": data["references"],
330 + "discovered_at": data["discovered_at"],
331 + "epss_score": data["epss_score"],
332 + "package_name": data["package_name"],
333 + },
334 + )
335 + )
336 + await db_session.execute(stmt)
337 + updated_count = len(update_data)
338 + logger.info(f"BULK MODE: Executed {updated_count} vulnerability updates")
339 +
340 + # Single commit for all operations
341 + await db_session.commit()
342 +
343 + total_synced = inserted_count + updated_count
344 + logger.info(f"BULK MODE: Successfully synced {total_synced} vulnerabilities ({inserted_count} new, {updated_count} updated)")
345 +
346 + return VulnerabilitySyncResponse(
347 + success=True,
348 + message=f"BULK MODE: Successfully synced {total_synced} vulnerabilities for agent {agent_name} ({inserted_count} new, {updated_count} updated)",
349 + synced_count=total_synced,
350 + errors=errors,
351 + )
352 +
353 + except Exception as e:
354 + await db_session.rollback()
355 + logger.error(f"BULK MODE: Error syncing vulnerabilities for agent {agent_name}: {e}")
356 + return VulnerabilitySyncResponse(
357 + success=False,
358 + message=f"BULK MODE: Failed to sync vulnerabilities for agent {agent_name}: {e}",
359 + synced_count=0,
360 + errors=[str(e)],
361 + )
362 +
363 +
364 +async def sync_vulnerabilities_for_agent(
365 + db_session: AsyncSession,
366 + agent_name: str,
367 + customer_code: Optional[str] = None,
368 + batch_size: int = 100,
369 + use_bulk_mode: bool = False,
370 +) -> VulnerabilitySyncResponse:
371 + """
372 + Sync vulnerabilities for a specific agent
373 +
374 + Args:
375 + db_session: Database session to use
376 + agent_name: Name of the agent to sync vulnerabilities for
377 + customer_code: Optional customer code override
378 + batch_size: Number of vulnerabilities to process in each batch (default: 100)
379 + use_bulk_mode: If True, use ultra-fast bulk operations (default: False)
380 +
381 + Returns:
382 + VulnerabilitySyncResponse with sync results
383 + """
384 + try:
385 + # Get agent from database using the session
386 + result = await db_session.execute(select(Agents).filter(Agents.hostname == agent_name))
387 + agent = result.scalars().first()
388 +
389 + if not agent:
390 + return VulnerabilitySyncResponse(
391 + success=False,
392 + message=f"Agent {agent_name} not found in database",
393 + synced_count=0,
394 + errors=[f"Agent {agent_name} not found"],
395 + )
396 +
397 + # Use agent's customer code if not provided
398 + if not customer_code:
399 + customer_code = agent.customer_code
400 +
401 + # Cache agent values to prevent lazy loading issues in the loop
402 + agent_id = agent.agent_id
403 +
404 + # Fetch vulnerabilities from Indexer
405 + vulnerability_docs = await fetch_vulnerabilities_from_indexer(agent_name=agent_name)
406 +
407 + logger.info(f"Fetched {len(vulnerability_docs)} vulnerabilities for agent {agent_name}")
408 +
409 + if not vulnerability_docs:
410 + return VulnerabilitySyncResponse(
411 + success=True,
412 + message=f"No vulnerabilities found for agent {agent_name}",
413 + synced_count=0,
414 + errors=[],
415 + )
416 +
417 + synced_count = 0
418 + errors = []
419 +
420 + # Choose processing mode based on use_bulk_mode flag
421 + if use_bulk_mode:
422 + logger.info(f"Using BULK MODE for {len(vulnerability_docs)} vulnerabilities for agent {agent_name}")
423 + return await _sync_vulnerabilities_bulk_mode(db_session, agent_id, agent_name, customer_code, vulnerability_docs)
424 +
425 + # OPTIMIZATION: Process vulnerabilities in batches for better performance
426 + logger.info(f"Using BATCH MODE (batch_size={batch_size}) for {len(vulnerability_docs)} vulnerabilities for agent {agent_name}")
427 +
428 + # First, get all existing vulnerabilities for this agent to do bulk comparison
429 + logger.info(f"Fetching existing vulnerabilities for agent {agent_name} for comparison")
430 + existing_vulns_result = await db_session.execute(select(AgentVulnerabilities).filter(AgentVulnerabilities.agent_id == agent_id))
431 + existing_vulns = existing_vulns_result.scalars().all()
432 +
433 + # Create a lookup dictionary for fast comparison (agent_id + cve_id + package_name)
434 + existing_vulns_lookup = {}
435 + for vuln in existing_vulns:
436 + key = f"{vuln.agent_id}_{vuln.cve_id}_{vuln.package_name or 'None'}"
437 + existing_vulns_lookup[key] = vuln
438 +
439 + logger.info(f"Found {len(existing_vulns_lookup)} existing vulnerabilities for agent {agent_name}")
440 +
441 + # Process vulnerabilities in batches
442 + for batch_start in range(0, len(vulnerability_docs), batch_size):
443 + batch_end = min(batch_start + batch_size, len(vulnerability_docs))
444 + batch_docs = vulnerability_docs[batch_start:batch_end]
445 +
446 + logger.info(
447 + f"Processing batch {batch_start // batch_size + 1}: vulnerabilities {batch_start + 1}-{batch_end} of {len(vulnerability_docs)} for agent {agent_name}",
448 + )
449 +
450 + try:
451 + batch_updates = []
452 + batch_inserts = []
453 + batch_errors = []
454 +
455 + # Process each document in the batch
456 + for i, doc in enumerate(batch_docs):
457 + try:
458 + # Process the vulnerability document
459 + vuln_data = process_wazuh_document(doc)
460 +
461 + # Create lookup key
462 + lookup_key = f"{agent_id}_{vuln_data.cve_id}_{vuln_data.package_name or 'None'}"
463 +
464 + if lookup_key in existing_vulns_lookup:
465 + # Update existing vulnerability
466 + existing_vuln = existing_vulns_lookup[lookup_key]
467 + existing_vuln.severity = vuln_data.severity
468 + existing_vuln.title = vuln_data.title
469 + existing_vuln.references = vuln_data.references
470 + existing_vuln.discovered_at = vuln_data.detected_at
471 + if hasattr(vuln_data, "base_score") and vuln_data.base_score:
472 + existing_vuln.epss_score = str(vuln_data.base_score)
473 + if hasattr(vuln_data, "package_name"):
474 + existing_vuln.package_name = vuln_data.package_name
475 +
476 + db_session.add(existing_vuln)
477 + batch_updates.append(vuln_data.cve_id)
478 + else:
479 + # Create new vulnerability record
480 + new_vuln = AgentVulnerabilities.create_from_model(
481 + vulnerability_data=vuln_data,
482 + agent_id=agent_id,
483 + customer_code=customer_code,
484 + )
485 + db_session.add(new_vuln)
486 + batch_inserts.append(vuln_data.cve_id)
487 +
488 + # Add to lookup to avoid duplicates within the same batch
489 + existing_vulns_lookup[lookup_key] = new_vuln
490 +
491 + except Exception as doc_error:
492 + error_msg = f"Error processing vulnerability {doc.get('_id', 'unknown')}: {doc_error}"
493 + logger.error(error_msg)
494 + batch_errors.append(error_msg)
495 + continue
496 +
497 + # Commit the entire batch at once
498 + await db_session.commit()
499 +
500 + batch_synced = len(batch_updates) + len(batch_inserts)
501 + synced_count += batch_synced
502 + errors.extend(batch_errors)
503 +
504 + logger.info(
505 + f"Batch {batch_start // batch_size + 1} completed: {len(batch_updates)} updates, {len(batch_inserts)} inserts, {len(batch_errors)} errors",
506 + )
507 +
508 + except Exception as batch_error:
509 + await db_session.rollback()
510 + error_msg = f"Error processing batch {batch_start}-{batch_end}: {batch_error}"
511 + logger.error(error_msg)
512 + errors.append(error_msg)
513 + continue
514 +
515 + return VulnerabilitySyncResponse(
516 + success=True,
517 + message=f"Successfully synced {synced_count} vulnerabilities for agent {agent_name} ({len(errors)}",
518 + synced_count=synced_count,
519 + errors=errors,
520 + )
521 +
522 + except Exception as e:
523 + await db_session.rollback()
524 + logger.error(f"Error syncing vulnerabilities for agent {agent_name}: {e}")
525 + return VulnerabilitySyncResponse(
526 + success=False,
527 + message=f"Failed to sync vulnerabilities for agent {agent_name}: {e}",
528 + synced_count=0,
529 + errors=[str(e)],
530 + )
531 +
532 +
533 +async def sync_all_vulnerabilities(
534 + db_session: AsyncSession,
535 + customer_code: Optional[str] = None,
536 + batch_size: int = 100,
537 + use_bulk_mode: bool = False,
538 +) -> VulnerabilitySyncResponse:
539 + """
540 + Sync vulnerabilities for all agents or agents of a specific customer with performance options
541 +
542 + Args:
543 + db_session: Database session to use
544 + customer_code: Optional customer code to filter agents by.
545 + If None, syncs vulnerabilities for all agents in database.
546 + batch_size: Number of vulnerabilities to process in each batch (default: 100)
547 + use_bulk_mode: Use ultra-fast bulk operations for large datasets (default: False)
548 +
549 + Returns:
550 + VulnerabilitySyncResponse with sync results
551 + """
552 + try:
553 + mode_info = "bulk mode" if use_bulk_mode else f"batch mode (size: {batch_size})"
554 + logger.info(f"Starting bulk vulnerability sync for customer: {customer_code or 'all agents'} using {mode_info}")
555 +
556 + # Build query to get agents using the session
557 + if customer_code:
558 + query = select(Agents).filter(Agents.customer_code == customer_code)
559 + else:
560 + query = select(Agents)
561 +
562 + # Execute query using the session
563 + result = await db_session.execute(query)
564 + agents = result.scalars().all()
565 +
566 + if not agents:
567 + message = "No agents found" + (f" for customer {customer_code}" if customer_code else " in database")
568 + return VulnerabilitySyncResponse(success=True, message=message, synced_count=0, errors=[])
569 +
570 + total_synced = 0
571 + all_errors = []
572 +
573 + # Process each agent synchronously to maintain session consistency
574 + for agent in agents:
575 + # Cache agent values to prevent lazy loading issues
576 + agent_hostname = agent.hostname
577 + agent_customer_code = agent.customer_code
578 +
579 + if not agent_hostname:
580 + continue
581 +
582 + try:
583 + logger.info(f"Starting sync for agent: {agent_hostname} using {mode_info}")
584 + result = await sync_vulnerabilities_for_agent(
585 + db_session=db_session,
586 + agent_name=agent_hostname,
587 + customer_code=agent_customer_code,
588 + batch_size=batch_size,
589 + use_bulk_mode=use_bulk_mode,
590 + )
591 +
592 + total_synced += result.synced_count
593 + all_errors.extend(result.errors)
594 +
595 + except Exception as agent_error:
596 + error_msg = f"Error syncing agent {agent_hostname}: {agent_error}"
597 + logger.error(error_msg)
598 + all_errors.append(error_msg)
599 +
600 + success_message = f"Completed vulnerability sync for {len(agents)} agents using {mode_info}"
601 + if customer_code:
602 + success_message += f" (customer: {customer_code})"
603 + else:
604 + success_message += " (all agents in database)"
605 +
606 + return VulnerabilitySyncResponse(success=True, message=success_message, synced_count=total_synced, errors=all_errors)
607 +
608 + except Exception as e:
609 + logger.error(f"Error in bulk vulnerability sync: {e}")
610 + return VulnerabilitySyncResponse(success=False, message=f"Failed to sync vulnerabilities: {e}", synced_count=0, errors=[str(e)])
611 +
612 +
613 +async def get_vulnerabilities_by_agent(
614 + db_session: AsyncSession,
615 + agent_id: str,
616 + severity_filter: Optional[List[str]] = None,
617 +) -> AgentVulnerabilitiesResponse:
618 + """
619 + Get vulnerabilities for a specific agent from database
620 +
621 + Args:
622 + db_session: Database session to use
623 + agent_id: Agent ID to get vulnerabilities for
624 + severity_filter: Optional list of severities to filter by
625 +
626 + Returns:
627 + AgentVulnerabilitiesResponse with vulnerabilities
628 + """
629 + try:
630 + query = select(AgentVulnerabilities).filter(AgentVulnerabilities.agent_id == agent_id)
631 +
632 + if severity_filter:
633 + query = query.filter(AgentVulnerabilities.severity.in_(severity_filter))
634 +
635 + result = await db_session.execute(query)
636 + vulnerabilities = result.scalars().all()
637 +
638 + vuln_list = [
639 + AgentVulnerabilityOut(
640 + id=vuln.id,
641 + cve_id=vuln.cve_id,
642 + severity=vuln.severity,
643 + title=vuln.title,
644 + references=vuln.references,
645 + status=vuln.status,
646 + discovered_at=vuln.discovered_at,
647 + remediated_at=vuln.remediated_at,
648 + epss_score=vuln.epss_score,
649 + epss_percentile=vuln.epss_percentile,
650 + package_name=vuln.package_name,
651 + agent_id=vuln.agent_id,
652 + customer_code=vuln.customer_code,
653 + )
654 + for vuln in vulnerabilities
655 + ]
656 +
657 + return AgentVulnerabilitiesResponse(
658 + vulnerabilities=vuln_list,
659 + success=True,
660 + message=f"Retrieved {len(vuln_list)} vulnerabilities for agent {agent_id}",
661 + total_count=len(vuln_list),
662 + )
663 +
664 + except Exception as e:
665 + logger.error(f"Error getting vulnerabilities for agent {agent_id}: {e}")
666 + raise HTTPException(status_code=500, detail=f"Failed to get vulnerabilities for agent {agent_id}: {e}")
667 +
668 +
669 +async def get_vulnerability_statistics(db_session: AsyncSession, customer_code: Optional[str] = None) -> VulnerabilityStatsResponse:
670 + """
671 + Get vulnerability statistics
672 +
673 + Args:
674 + db_session: Database session to use
675 + customer_code: Optional customer code to filter by
676 +
677 + Returns:
678 + VulnerabilityStatsResponse with statistics
679 + """
680 + try:
681 + query = select(AgentVulnerabilities)
682 + if customer_code:
683 + query = query.filter(AgentVulnerabilities.customer_code == customer_code)
684 +
685 + result = await db_session.execute(query)
686 + vulnerabilities = result.scalars().all()
687 +
688 + # Calculate statistics
689 + total = len(vulnerabilities)
690 + critical = sum(1 for v in vulnerabilities if v.severity.lower() == "critical")
691 + high = sum(1 for v in vulnerabilities if v.severity.lower() == "high")
692 + medium = sum(1 for v in vulnerabilities if v.severity.lower() == "medium")
693 + low = sum(1 for v in vulnerabilities if v.severity.lower() == "low")
694 +
695 + # Group by customer if no specific customer requested
696 + by_customer = {}
697 + if not customer_code:
698 + for vuln in vulnerabilities:
699 + if vuln.customer_code:
700 + by_customer[vuln.customer_code] = by_customer.get(vuln.customer_code, 0) + 1
701 +
702 + return VulnerabilityStatsResponse(
703 + total_vulnerabilities=total,
704 + critical_count=critical,
705 + high_count=high,
706 + medium_count=medium,
707 + low_count=low,
708 + by_customer=by_customer,
709 + success=True,
710 + message="Vulnerability statistics retrieved successfully",
711 + )
712 +
713 + except Exception as e:
714 + logger.error(f"Error getting vulnerability statistics: {e}")
715 + raise HTTPException(status_code=500, detail=f"Failed to get vulnerability statistics: {e}")
716 +
717 +
718 +async def delete_vulnerabilities(db_session: AsyncSession, agent_name: Optional[str] = None, customer_code: Optional[str] = None):
719 + """
720 + Delete vulnerabilities based on scope:
721 + - If neither agent_name nor customer_code provided: Delete ALL vulnerabilities
722 + - If agent_name provided: Delete vulnerabilities for that specific agent
723 + - If customer_code provided: Delete vulnerabilities for all agents of that customer
724 +
725 + Args:
726 + db_session: Database session to use
727 + agent_name: Optional agent name to delete vulnerabilities for
728 + customer_code: Optional customer code to delete vulnerabilities for
729 +
730 + Returns:
731 + VulnerabilityDeleteResponse with deletion results
732 + """
733 + from app.agents.vulnerabilities.schema.vulnerabilities import (
734 + VulnerabilityDeleteResponse,
735 + )
736 +
737 + try:
738 + deleted_count = 0
739 +
740 + if agent_name:
741 + # Delete vulnerabilities for specific agent
742 + logger.info(f"Deleting vulnerabilities for agent: {agent_name}")
743 +
744 + # First get the agent to validate it exists and get agent_id
745 + agent_result = await db_session.execute(select(Agents).filter(Agents.hostname == agent_name))
746 + agent = agent_result.scalars().first()
747 +
748 + if not agent:
749 + return VulnerabilityDeleteResponse(
750 + success=False,
751 + message=f"Agent {agent_name} not found in database",
752 + deleted_count=0,
753 + errors=[f"Agent {agent_name} not found"],
754 + )
755 +
756 + # Delete vulnerabilities for this agent
757 + from sqlalchemy import delete
758 +
759 + delete_stmt = delete(AgentVulnerabilities).where(AgentVulnerabilities.agent_id == agent.agent_id)
760 + result = await db_session.execute(delete_stmt)
761 + deleted_count = result.rowcount
762 + await db_session.commit()
763 +
764 + return VulnerabilityDeleteResponse(
765 + success=True,
766 + message=f"Successfully deleted {deleted_count} vulnerabilities for agent {agent_name}",
767 + deleted_count=deleted_count,
768 + errors=[],
769 + )
770 +
771 + elif customer_code:
772 + # Delete vulnerabilities for all agents of specific customer
773 + logger.info(f"Deleting vulnerabilities for customer: {customer_code}")
774 +
775 + from sqlalchemy import delete
776 +
777 + delete_stmt = delete(AgentVulnerabilities).where(AgentVulnerabilities.customer_code == customer_code)
778 + result = await db_session.execute(delete_stmt)
779 + deleted_count = result.rowcount
780 + await db_session.commit()
781 +
782 + return VulnerabilityDeleteResponse(
783 + success=True,
784 + message=f"Successfully deleted {deleted_count} vulnerabilities for customer {customer_code}",
785 + deleted_count=deleted_count,
786 + errors=[],
787 + )
788 +
789 + else:
790 + # Delete ALL vulnerabilities
791 + logger.warning("Deleting ALL vulnerabilities from database")
792 +
793 + from sqlalchemy import delete
794 +
795 + delete_stmt = delete(AgentVulnerabilities)
796 + result = await db_session.execute(delete_stmt)
797 + deleted_count = result.rowcount
798 + await db_session.commit()
799 +
800 + return VulnerabilityDeleteResponse(
801 + success=True,
802 + message=f"Successfully deleted ALL {deleted_count} vulnerabilities from database",
803 + deleted_count=deleted_count,
804 + errors=[],
805 + )
806 +
807 + except Exception as e:
808 + await db_session.rollback()
809 + logger.error(f"Error deleting vulnerabilities: {e}")
810 + return VulnerabilityDeleteResponse(
811 + success=False,
812 + message=f"Failed to delete vulnerabilities: {e}",
813 + deleted_count=0,
814 + errors=[str(e)],
815 + )
816 +
817 +
818 +async def search_vulnerabilities_from_indexer(
819 + db_session: AsyncSession,
820 + customer_code: Optional[str] = None,
821 + agent_name: Optional[str] = None,
822 + severity: Optional[str] = None,
823 + cve_id: Optional[str] = None,
824 + package_name: Optional[str] = None,
825 + page: int = 1,
826 + page_size: int = 50,
827 + include_epss: bool = True,
828 +) -> VulnerabilitySearchResponse:
829 + """
830 + Search vulnerabilities directly from Wazuh indexer with filtering and pagination
831 +
832 + Args:
833 + db_session: Database session for agent lookup
834 + customer_code: Optional customer code filter
835 + agent_name: Optional agent hostname filter
836 + severity: Optional severity filter
837 + cve_id: Optional CVE ID filter
838 + package_name: Optional package name filter
839 + page: Page number for pagination
840 + page_size: Number of results per page
841 + include_epss: Whether to include EPSS scores (default: True, may impact performance)
842 +
843 + Returns:
844 + VulnerabilitySearchResponse with paginated results
845 + """
846 + logger.info(
847 + f"Searching vulnerabilities with filters: customer_code={customer_code}, "
848 + f"agent_name={agent_name}, severity={severity}, cve_id={cve_id}, "
849 + f"package_name={package_name}, page={page}, page_size={page_size}, "
850 + f"include_epss={include_epss}",
851 + )
852 +
853 + # Build filters applied dict for response
854 + filters_applied = {}
855 + if customer_code:
856 + filters_applied["customer_code"] = customer_code
857 + if agent_name:
858 + filters_applied["agent_name"] = agent_name
859 + if severity:
860 + filters_applied["severity"] = severity
861 + if cve_id:
862 + filters_applied["cve_id"] = cve_id
863 + if package_name:
864 + filters_applied["package_name"] = package_name
865 +
866 + # Create Elasticsearch client
867 + es_client = None
868 + try:
869 + # Initialize Elasticsearch client
870 + es_client = await create_wazuh_indexer_client_async("Wazuh-Indexer")
871 +
872 + # Always get all agents to build hostname to customer_code mapping
873 + # This ensures we can always provide customer_code in the response
874 + all_agents_query = select(Agents)
875 + all_agents_result = await db_session.execute(all_agents_query)
876 + all_agents = all_agents_result.scalars().all()
877 +
878 + # Build complete agent hostname to customer code mapping
879 + customer_agent_map = {}
880 + for agent in all_agents:
881 + if agent.hostname:
882 + customer_agent_map[agent.hostname] = agent.customer_code
883 +
884 + # Get agent information for filtering (if filters are applied)
885 + agent_hostnames = []
886 +
887 + if customer_code or agent_name:
888 + query = select(Agents)
889 + if customer_code:
890 + query = query.filter(Agents.customer_code == customer_code)
891 + if agent_name:
892 + query = query.filter(Agents.hostname == agent_name)
893 +
894 + result = await db_session.execute(query)
895 + agents = result.scalars().all()
896 +
897 + if not agents and (customer_code or agent_name):
898 + return VulnerabilitySearchResponse(
899 + vulnerabilities=[],
900 + total_count=0,
901 + page=page,
902 + page_size=page_size,
903 + total_pages=0,
904 + has_next=False,
905 + has_previous=False,
906 + success=True,
907 + message="No agents found matching the specified criteria",
908 + filters_applied=filters_applied,
909 + )
910 +
911 + # Build list of agent hostnames for Elasticsearch filtering
912 + for agent in agents:
913 + if agent.hostname:
914 + agent_hostnames.append(agent.hostname)
915 +
916 + # Get vulnerability indices
917 + vuln_indices = await get_vulnerabilities_indices()
918 + if not vuln_indices:
919 + return VulnerabilitySearchResponse(
920 + vulnerabilities=[],
921 + total_count=0,
922 + page=page,
923 + page_size=page_size,
924 + total_pages=0,
925 + has_next=False,
926 + has_previous=False,
927 + success=True,
928 + message="No vulnerability indices found",
929 + filters_applied=filters_applied,
930 + )
931 +
932 + # Build Elasticsearch query
933 + es_query = {"bool": {"must": []}}
934 +
935 + # Add agent filter if specified
936 + if agent_hostnames:
937 + es_query["bool"]["must"].append({"terms": {"agent.name": agent_hostnames}})
938 +
939 + # Add severity filter
940 + if severity:
941 + es_query["bool"]["must"].append({"term": {"vulnerability.severity": severity}})
942 +
943 + # Add CVE ID filter
944 + if cve_id:
945 + es_query["bool"]["must"].append({"term": {"vulnerability.id": cve_id}})
946 +
947 + # Add package name filter
948 + if package_name:
949 + es_query["bool"]["must"].append({"wildcard": {"package.name": f"*{package_name}*"}})
950 +
951 + # Calculate pagination
952 + start_index = (page - 1) * page_size
953 +
954 + # Create Elasticsearch client
955 + es_client = None
956 + try:
957 + es_client = await create_wazuh_indexer_client_async("Wazuh-Indexer")
958 +
959 + # First, get total count
960 + count_response = await es_client.count(index=",".join(vuln_indices), body={"query": es_query})
961 + total_count = count_response["count"]
962 +
963 + # Calculate pagination info
964 + total_pages = (total_count + page_size - 1) // page_size
965 + has_next = page < total_pages
966 + has_previous = page > 1
967 +
968 + if total_count == 0:
969 + return VulnerabilitySearchResponse(
970 + vulnerabilities=[],
971 + total_count=0,
972 + page=page,
973 + page_size=page_size,
974 + total_pages=0,
975 + has_next=False,
976 + has_previous=False,
977 + success=True,
978 + message="No vulnerabilities found matching the specified criteria",
979 + filters_applied=filters_applied,
980 + )
981 +
982 + # Get the actual results with pagination
983 + search_response = await es_client.search(
984 + index=",".join(vuln_indices),
985 + body={
986 + "query": es_query,
987 + "sort": [{"vulnerability.detected_at": {"order": "desc"}}, {"vulnerability.severity": {"order": "asc"}}],
988 + "from": start_index,
989 + "size": page_size,
990 + },
991 + )
992 +
993 + vulnerabilities = []
994 + for hit in search_response["hits"]["hits"]:
995 + try:
996 + source = hit["_source"]
997 + agent_data = source.get("agent", {})
998 + agent_hostname = agent_data.get("name", "unknown")
999 +
1000 + # Get customer code from our mapping
1001 + agent_customer_code = customer_agent_map.get(agent_hostname)
1002 +
1003 + # Process the vulnerability data
1004 + vuln_data = process_wazuh_document(hit)
1005 +
1006 + # Get EPSS score for the CVE (if requested)
1007 + epss_score, epss_percentile = None, None
1008 + if include_epss:
1009 + epss_score, epss_percentile = await get_epss_score_for_cve(vuln_data.cve_id)
1010 +
1011 + vulnerability_item = VulnerabilitySearchItem(
1012 + cve_id=vuln_data.cve_id,
1013 + severity=vuln_data.severity,
1014 + title=vuln_data.title,
1015 + agent_name=agent_hostname,
1016 + customer_code=agent_customer_code,
1017 + references=vuln_data.references,
1018 + detected_at=vuln_data.detected_at,
1019 + published_at=vuln_data.published_at,
1020 + base_score=vuln_data.base_score,
1021 + package_name=vuln_data.package_name,
1022 + package_version=vuln_data.package_version,
1023 + package_architecture=vuln_data.package_architecture,
1024 + epss_score=epss_score,
1025 + epss_percentile=epss_percentile,
1026 + )
1027 + vulnerabilities.append(vulnerability_item)
1028 +
1029 + except Exception as e:
1030 + logger.error(f"Error processing vulnerability document: {e}")
1031 + continue
1032 +
1033 + message = f"Found {len(vulnerabilities)} vulnerabilities on page {page} of {total_pages}"
1034 + if filters_applied:
1035 + message += f" with filters: {filters_applied}"
1036 +
1037 + return VulnerabilitySearchResponse(
1038 + vulnerabilities=vulnerabilities,
1039 + total_count=total_count,
1040 + page=page,
1041 + page_size=page_size,
1042 + total_pages=total_pages,
1043 + has_next=has_next,
1044 + has_previous=has_previous,
1045 + success=True,
1046 + message=message,
1047 + filters_applied=filters_applied,
1048 + )
1049 +
1050 + except Exception as e:
1051 + logger.error(f"Error searching vulnerabilities from indexer: {e}")
1052 + return VulnerabilitySearchResponse(
1053 + vulnerabilities=[],
1054 + total_count=0,
1055 + page=page,
1056 + page_size=page_size,
1057 + total_pages=0,
1058 + has_next=False,
1059 + has_previous=False,
1060 + success=False,
1061 + message=f"Failed to search vulnerabilities: {e}",
1062 + filters_applied=filters_applied if "filters_applied" in locals() else {},
1063 + )
1064 + finally:
1065 + # Ensure the Elasticsearch client session is properly closed
1066 + if es_client:
1067 + try:
1068 + await es_client.close()
1069 + except Exception as close_error:
1070 + logger.warning(f"Error closing Elasticsearch client: {close_error}")
1071 +
1072 + except Exception as e:
1073 + logger.error(f"Unexpected error in search_vulnerabilities_from_indexer: {e}")
1074 + return VulnerabilitySearchResponse(
1075 + vulnerabilities=[],
1076 + total_count=0,
1077 + page=page,
1078 + page_size=page_size,
1079 + total_pages=0,
1080 + has_next=False,
1081 + has_previous=False,
1082 + success=False,
1083 + message=f"Unexpected error occurred: {e}",
1084 + filters_applied=filters_applied if "filters_applied" in locals() else {},
1085 + )
backend/app/db/universal_models.py
+64
@@ -106,6 +106,7 @@ class Agents(SQLModel, table=True):
106 velociraptor_org: Optional[str] = Field(max_length=256)
107
108 customer: Optional[Customers] = Relationship(back_populates="agents")
109 + vulnerabilities: Optional[list["AgentVulnerabilities"]] = Relationship(back_populates="agent")
110
111 @classmethod
112 def create_from_model(cls, wazuh_agent, velociraptor_agent, customer_code):
@@ -257,3 +258,66 @@ class SchedulerJob(SQLModel, table=True):
258
259 def __repr__(self):
260 return f"<SchedulerJob(id={self.id}, next_run_time={self.next_run_time})>"
261 +
262 +
263 +class AgentVulnerabilities(SQLModel, table=True):
264 + __tablename__ = "agent_vulnerabilities"
265 +
266 + id: Optional[int] = Field(primary_key=True)
267 + cve_id: str = Field(default="UNKNOWN_CVE", max_length=50, index=True)
268 + severity: str = Field(default="UNKNOWN", max_length=50, index=True)
269 + title: str = Field(max_length=255)
270 + references: str = Field(default=None, max_length=2048)
271 + status: str = Field(default="Active", max_length=50, index=True)
272 + discovered_at: datetime = Field(index=True)
273 + remediated_at: Optional[datetime] = Field(default=None)
274 + epss_score: Optional[str] = Field(default=None, max_length=50)
275 + epss_percentile: Optional[str] = Field(default=None, max_length=50)
276 + package_name: Optional[str] = Field(default=None, max_length=255)
277 +
278 + # Foreign keys
279 + agent_id: str = Field(foreign_key="agents.agent_id", max_length=256, index=True)
280 + customer_code: Optional[str] = Field(foreign_key="customers.customer_code", max_length=50, index=True)
281 +
282 + # Relationship back to the Agents model
283 + agent: Optional["Agents"] = Relationship(back_populates="vulnerabilities")
284 +
285 + def update_from_model(self, vulnerability_data):
286 + """Update vulnerability from external data source"""
287 + if hasattr(vulnerability_data, "cve_id"):
288 + self.cve_id = vulnerability_data.cve_id
289 + if hasattr(vulnerability_data, "severity"):
290 + self.severity = vulnerability_data.severity
291 + if hasattr(vulnerability_data, "title"):
292 + self.title = vulnerability_data.title
293 + if hasattr(vulnerability_data, "references"):
294 + self.references = vulnerability_data.references
295 + if hasattr(vulnerability_data, "detected_at"):
296 + self.discovered_at = vulnerability_data.detected_at
297 + if hasattr(vulnerability_data, "status"):
298 + self.status = vulnerability_data.status
299 + if hasattr(vulnerability_data, "epss_score"):
300 + self.epss_score = vulnerability_data.epss_score
301 + if hasattr(vulnerability_data, "epss_percentile"):
302 + self.epss_percentile = vulnerability_data.epss_percentile
303 + if hasattr(vulnerability_data, "package_name"):
304 + self.package_name = vulnerability_data.package_name
305 + if hasattr(vulnerability_data, "remediated_at"):
306 + self.remediated_at = vulnerability_data.remediated_at
307 +
308 + @classmethod
309 + def create_from_model(cls, vulnerability_data, agent_id, customer_code=None):
310 + """Create a new vulnerability record from external data"""
311 + return cls(
312 + cve_id=getattr(vulnerability_data, "cve_id", "UNKNOWN_CVE"),
313 + severity=getattr(vulnerability_data, "severity", "UNKNOWN"),
314 + title=getattr(vulnerability_data, "title", ""),
315 + references=getattr(vulnerability_data, "references", None),
316 + status=getattr(vulnerability_data, "status", "Active"),
317 + epss_score=getattr(vulnerability_data, "epss_score", None),
318 + epss_percentile=getattr(vulnerability_data, "epss_percentile", None),
319 + package_name=getattr(vulnerability_data, "package_name", None),
320 + discovered_at=getattr(vulnerability_data, "detected_at", datetime.utcnow()),
321 + agent_id=agent_id,
322 + customer_code=customer_code,
323 + )
backend/app/routers/agents.py
+2
@@ -1,9 +1,11 @@
1 from fastapi import APIRouter
2
3 from app.agents.routes.agents import agents_router
4 +from app.agents.vulnerabilities.routes.vulnerabilities import vulnerabilities_router
5
6 # Instantiate the APIRouter
7 router = APIRouter()
8
9 # Include the Wazuh Manager related routes
10 router.include_router(agents_router, prefix="/agents", tags=["agents"])
11 +router.include_router(vulnerabilities_router, prefix="/vulnerabilities", tags=["vulnerabilities"])