| 1 | import asyncio |
| 2 | from typing import List |
| 3 | |
| 4 | from fastapi import HTTPException |
| 5 | from loguru import logger |
| 6 | |
| 7 | from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilities |
| 8 | from app.agents.wazuh.schema.agents import WazuhAgentVulnerabilitiesResponse |
| 9 | from app.connectors.wazuh_indexer.utils.universal import collect_indices |
| 10 | from app.connectors.wazuh_indexer.utils.universal import create_wazuh_indexer_client |
| 11 | from app.connectors.wazuh_indexer.utils.universal import ( |
| 12 | create_wazuh_indexer_client_async, |
| 13 | ) |
| 14 | from app.connectors.wazuh_manager.utils.universal import send_get_request |
| 15 | from app.integrations.utils.event_shipper import event_shipper |
| 16 | from app.integrations.utils.schema import EventShipperPayload |
| 17 | |
| 18 | |
| 19 | async def collect_agent_vulnerabilities(agent_id: str, vulnerability_severity: str): |
| 20 | """ |
| 21 | Collect agent vulnerabilities from Wazuh Manager. |
| 22 | Used when Wazuh Manager is below 4.8.0 |
| 23 | |
| 24 | Args: |
| 25 | agent_id (str): The ID of the agent. |
| 26 | vulnerability_severity (str): The severity of the vulnerabilities to collect. |
| 27 | |
| 28 | Returns: |
| 29 | WazuhAgentVulnerabilitiesResponse: An object containing the collected vulnerabilities. |
| 30 | |
| 31 | Raises: |
| 32 | HTTPException: If there is an error collecting the vulnerabilities. |
| 33 | """ |
| 34 | logger.info(f"Collecting agent {agent_id} vulnerabilities from Wazuh Manager") |
| 35 | |
| 36 | severities = ["Low", "Medium", "High", "Critical"] if vulnerability_severity == "All" else [vulnerability_severity] |
| 37 | |
| 38 | agent_vulnerabilities = [] |
| 39 | for severity in severities: |
| 40 | response = await send_get_request( |
| 41 | endpoint=f"/vulnerability/{agent_id}", |
| 42 | params={"severity": severity}, |
| 43 | ) |
| 44 | if response["success"] is False: |
| 45 | raise HTTPException(status_code=500, detail=response["message"]) |
| 46 | # Navigate through the nested 'data' structure to get 'affected_items' |
| 47 | affected_items = response.get("data", {}).get("data", {}).get("affected_items", []) |
| 48 | agent_vulnerabilities.extend(affected_items) |
| 49 | |
| 50 | processed_vulnerabilities = process_agent_vulnerabilities(agent_vulnerabilities) |
| 51 | |
| 52 | return WazuhAgentVulnerabilitiesResponse( |
| 53 | vulnerabilities=processed_vulnerabilities, |
| 54 | success=True, |
| 55 | message="Vulnerabilities collected successfully", |
| 56 | ) |
| 57 | |
| 58 | |
| 59 | def process_agent_vulnerabilities( |
| 60 | agent_vulnerabilities: List[dict], |
| 61 | ) -> List[WazuhAgentVulnerabilities]: |
| 62 | """ |
| 63 | Process agent vulnerabilities and return a list of WazuhAgentVulnerabilities objects. |
| 64 | |
| 65 | Args: |
| 66 | agent_vulnerabilities (List[dict]): A list of dictionaries containing agent vulnerabilities data. |
| 67 | |
| 68 | Returns: |
| 69 | List[WazuhAgentVulnerabilities]: A list of WazuhAgentVulnerabilities objects. |
| 70 | |
| 71 | Raises: |
| 72 | HTTPException: If there is an error processing the agent vulnerabilities. |
| 73 | """ |
| 74 | try: |
| 75 | return [WazuhAgentVulnerabilities(**vuln) for vuln in agent_vulnerabilities] |
| 76 | except Exception as e: |
| 77 | raise HTTPException( |
| 78 | status_code=500, |
| 79 | detail=f"Failed to process agent vulnerabilities: {e}", |
| 80 | ) |
| 81 | |
| 82 | |
| 83 | async def collect_agent_vulnerabilities_new(agent_id: str, vulnerability_severity: str): |
| 84 | """ |
| 85 | Collects vulnerabilities for a specific agent from the Wazuh Indexer Index. |
| 86 | Used when Wazuh-Manager is 4.8.0 or above. |
| 87 | |
| 88 | Args: |
| 89 | agent_id (str): The ID of the agent for which to collect vulnerabilities. |
| 90 | |
| 91 | Returns: |
| 92 | WazuhAgentVulnerabilitiesResponse: An object containing the collected vulnerabilities, |
| 93 | along with a success flag and a message indicating the success status. |
| 94 | """ |
| 95 | logger.info(f"Collecting agent {agent_id} vulnerabilities from Wazuh Indexer Index") |
| 96 | es = await create_wazuh_indexer_client("Wazuh-Indexer") |
| 97 | indices = await collect_indices(all_indices=True) |
| 98 | logger.info(f"Indices collect: {indices}") |
| 99 | |
| 100 | vulnerabilities_indices = filter_vulnerabilities_indices(indices.indices_list) |
| 101 | |
| 102 | agent_vulnerabilities = await collect_vulnerabilities(es, vulnerabilities_indices, agent_id, vulnerability_severity) |
| 103 | |
| 104 | processed_vulnerabilities = process_agent_vulnerabilities_new(agent_vulnerabilities) |
| 105 | |
| 106 | return WazuhAgentVulnerabilitiesResponse( |
| 107 | vulnerabilities=processed_vulnerabilities, |
| 108 | success=True, |
| 109 | message="Vulnerabilities collected successfully", |
| 110 | ) |
| 111 | |
| 112 | |
| 113 | def filter_vulnerabilities_indices(indices_list): |
| 114 | return [index for index in indices_list if index.startswith("wazuh-states-vulnerabilities")] |
| 115 | |
| 116 | |
| 117 | def filter_vulnerabilities_indices_sync(indices_list, customer_code): |
| 118 | """ |
| 119 | Filter the indices list to only include the vulnerability indices which are relevant to the customer. |
| 120 | Notice the missing `states` in the index name. |
| 121 | """ |
| 122 | # ! Make the customer code lowercase ! # |
| 123 | return [index for index in indices_list if index.startswith(f"wazuh-vulnerabilities-{customer_code.lower()}")] |
| 124 | |
| 125 | |
| 126 | async def collect_vulnerabilities(es, vulnerabilities_indices, agent_id, vulnerability_severity="Critical"): |
| 127 | agent_vulnerabilities = [] |
| 128 | for index in vulnerabilities_indices: |
| 129 | if vulnerability_severity == "All": |
| 130 | query = { |
| 131 | "query": { |
| 132 | "bool": { |
| 133 | "must": [ |
| 134 | {"match": {"agent.id": agent_id}}, |
| 135 | {"terms": {"vulnerability.severity": ["Low", "Medium", "High", "Critical"]}}, |
| 136 | ], |
| 137 | }, |
| 138 | }, |
| 139 | } |
| 140 | else: |
| 141 | query = { |
| 142 | "query": { |
| 143 | "bool": {"must": [{"match": {"agent.id": agent_id}}, {"match": {"vulnerability.severity": vulnerability_severity}}]}, |
| 144 | }, |
| 145 | } |
| 146 | |
| 147 | page = es.search(index=index, body=query, scroll="2m") |
| 148 | sid = page["_scroll_id"] |
| 149 | scroll_size = len(page["hits"]["hits"]) |
| 150 | |
| 151 | while scroll_size > 0: |
| 152 | for hit in page["hits"]["hits"]: |
| 153 | vulnerability = hit["_source"] |
| 154 | agent_vulnerabilities.append(vulnerability) |
| 155 | |
| 156 | page = es.scroll(scroll_id=sid, scroll="2m") |
| 157 | sid = page["_scroll_id"] |
| 158 | scroll_size = len(page["hits"]["hits"]) |
| 159 | |
| 160 | return agent_vulnerabilities |
| 161 | |
| 162 | |
| 163 | async def collect_vulnerabilities_sync(es, vulnerabilities_indices, agent_name, vulnerability_severity="All"): |
| 164 | agent_vulnerabilities = [] |
| 165 | for index in vulnerabilities_indices: |
| 166 | if vulnerability_severity == "All": |
| 167 | query = { |
| 168 | "query": { |
| 169 | "bool": { |
| 170 | "must": [ |
| 171 | {"match": {"agent.name": agent_name}}, |
| 172 | {"terms": {"vulnerability.severity": ["Low", "Medium", "High", "Critical"]}}, |
| 173 | ], |
| 174 | }, |
| 175 | }, |
| 176 | } |
| 177 | else: |
| 178 | query = { |
| 179 | "query": { |
| 180 | "bool": { |
| 181 | "must": [{"match": {"agent.name": agent_name}}, {"match": {"vulnerability.severity": vulnerability_severity}}], |
| 182 | }, |
| 183 | }, |
| 184 | } |
| 185 | |
| 186 | page = es.search(index=index, body=query, scroll="2m") |
| 187 | sid = page["_scroll_id"] |
| 188 | scroll_size = len(page["hits"]["hits"]) |
| 189 | |
| 190 | while scroll_size > 0: |
| 191 | for hit in page["hits"]["hits"]: |
| 192 | vulnerability = hit["_source"] |
| 193 | agent_vulnerabilities.append(vulnerability) |
| 194 | |
| 195 | page = es.scroll(scroll_id=sid, scroll="2m") |
| 196 | sid = page["_scroll_id"] |
| 197 | scroll_size = len(page["hits"]["hits"]) |
| 198 | |
| 199 | return agent_vulnerabilities |
| 200 | |
| 201 | |
| 202 | async def collect_vulnerabilities_async(es, vulnerabilities_indices, agent_name, vulnerability_severity="All"): |
| 203 | agent_vulnerabilities = [] |
| 204 | for index in vulnerabilities_indices: |
| 205 | if vulnerability_severity == "All": |
| 206 | query = { |
| 207 | "query": { |
| 208 | "bool": { |
| 209 | "must": [ |
| 210 | {"match": {"agent.name": agent_name}}, |
| 211 | {"terms": {"vulnerability.severity": ["Low", "Medium", "High", "Critical"]}}, |
| 212 | ], |
| 213 | }, |
| 214 | }, |
| 215 | } |
| 216 | else: |
| 217 | query = { |
| 218 | "query": { |
| 219 | "bool": { |
| 220 | "must": [{"match": {"agent.name": agent_name}}, {"match": {"vulnerability.severity": vulnerability_severity}}], |
| 221 | }, |
| 222 | }, |
| 223 | } |
| 224 | |
| 225 | page = await es.search(index=index, body=query, scroll="2m") |
| 226 | sid = page["_scroll_id"] |
| 227 | scroll_size = len(page["hits"]["hits"]) |
| 228 | |
| 229 | while scroll_size > 0: |
| 230 | for hit in page["hits"]["hits"]: |
| 231 | vulnerability = hit["_source"] |
| 232 | agent_vulnerabilities.append(vulnerability) |
| 233 | |
| 234 | page = await es.scroll(scroll_id=sid, scroll="2m") |
| 235 | sid = page["_scroll_id"] |
| 236 | scroll_size = len(page["hits"]["hits"]) |
| 237 | |
| 238 | return agent_vulnerabilities |
| 239 | |
| 240 | |
| 241 | def process_agent_vulnerabilities_new(agent_vulnerabilities: List[dict]) -> List[WazuhAgentVulnerabilities]: |
| 242 | logger.info(f"Processing agent vulnerabilities: {agent_vulnerabilities}") |
| 243 | |
| 244 | processed_vulnerabilities = [] |
| 245 | for vulnerability in agent_vulnerabilities: |
| 246 | processed_vulnerability = process_single_vulnerability(vulnerability) |
| 247 | processed_vulnerabilities.append(processed_vulnerability) |
| 248 | |
| 249 | return processed_vulnerabilities |
| 250 | |
| 251 | |
| 252 | def process_single_vulnerability(vulnerability): |
| 253 | external_references = ensure_list(vulnerability.get("vulnerability").get("reference")) |
| 254 | return WazuhAgentVulnerabilities( |
| 255 | severity=vulnerability.get("vulnerability").get("severity"), |
| 256 | version=vulnerability.get("package").get("version"), |
| 257 | type=vulnerability.get("package").get("type"), |
| 258 | name=vulnerability.get("package").get("name"), |
| 259 | external_references=external_references, |
| 260 | detection_time=vulnerability.get("vulnerability").get("detected_at"), |
| 261 | cvss3_score=vulnerability.get("vulnerability").get("score").get("base"), |
| 262 | published=vulnerability.get("vulnerability").get("published_at"), |
| 263 | architecture=vulnerability.get("package").get("architecture"), |
| 264 | cve=vulnerability.get("vulnerability").get("id"), |
| 265 | status=vulnerability.get("status"), |
| 266 | title=vulnerability.get("vulnerability").get("description"), |
| 267 | ) |
| 268 | |
| 269 | |
| 270 | def ensure_list(value): |
| 271 | if not isinstance(value, list): |
| 272 | return [value] |
| 273 | return value |
| 274 | |
| 275 | |
| 276 | async def check_vulnerability_exists_async(es, vulnerability_cve, agent_name, index_prefix): |
| 277 | query = { |
| 278 | "query": { |
| 279 | "bool": { |
| 280 | "must": [ |
| 281 | {"match": {"agent_name": agent_name}}, |
| 282 | {"match": {"cve": vulnerability_cve}}, |
| 283 | ], |
| 284 | }, |
| 285 | }, |
| 286 | } |
| 287 | index_pattern = f"{index_prefix}*" |
| 288 | response = await es.search(index=index_pattern, body=query) |
| 289 | return response["hits"]["total"]["value"] > 0 |
| 290 | |
| 291 | |
| 292 | async def sync_agent_vulnerabilities(agent_name: str, customer_code: str): |
| 293 | """ |
| 294 | 1. Loops through all agents in the database to collect their agent_name and customer code. |
| 295 | 2. Queries the `wazuh-states-vulnerabilities-*` index in Wazuh Indexer to get vulnerabilities based on the agent_name. |
| 296 | 3. Checks the `wazuh-vulnerabilities-*customer_code*` index in Wazuh Indexer to get vulnerabilities based on the |
| 297 | agent_name and checks to see if a vulnerability_id already exists. |
| 298 | 4. If the vulnerability_id does not exist, it is sent to the Graylog GELF Input. |
| 299 | """ |
| 300 | logger.info(f"Syncing agent {agent_name} with customer code {customer_code} vulnerabilities") |
| 301 | |
| 302 | es = await create_wazuh_indexer_client_async("Wazuh-Indexer") |
| 303 | indices = await collect_indices(all_indices=True) |
| 304 | |
| 305 | vulnerabilities_indices = filter_vulnerabilities_indices(indices.indices_list) |
| 306 | |
| 307 | agent_vulnerabilities = await collect_vulnerabilities_async(es, vulnerabilities_indices, agent_name, vulnerability_severity="All") |
| 308 | |
| 309 | processed_vulnerabilities = process_agent_vulnerabilities_new(agent_vulnerabilities) |
| 310 | |
| 311 | customer_vulnerabilities_indices = filter_vulnerabilities_indices_sync(indices.indices_list, customer_code) |
| 312 | |
| 313 | if customer_vulnerabilities_indices: |
| 314 | logger.info("Customer vulnerabilities index already exists") |
| 315 | # Create a list of tasks for checking vulnerabilities |
| 316 | tasks = [ |
| 317 | check_vulnerability_exists_async( |
| 318 | es, |
| 319 | vulnerability_cve=vulnerability.cve, |
| 320 | agent_name=agent_name, |
| 321 | index_prefix=f"wazuh-vulnerabilities-{customer_code}", |
| 322 | ) |
| 323 | for vulnerability in processed_vulnerabilities |
| 324 | ] |
| 325 | |
| 326 | # Run all tasks concurrently |
| 327 | results = await asyncio.gather(*tasks) |
| 328 | |
| 329 | # Process the results |
| 330 | for vulnerability, vulnerability_exists in zip(processed_vulnerabilities, results): |
| 331 | if not vulnerability_exists: |
| 332 | logger.info( |
| 333 | f"Vulnerability {vulnerability.cve} does not exist in customer index for agent {agent_name}, sending to Graylog", |
| 334 | ) |
| 335 | await event_shipper( |
| 336 | EventShipperPayload( |
| 337 | integration="vulnerabilities", |
| 338 | customer_code=customer_code, |
| 339 | agent_name=agent_name, |
| 340 | **vulnerability.model_dump(), |
| 341 | ), |
| 342 | ) |
| 343 | return True |
| 344 | |
| 345 | logger.info("Customer vulnerabilities index does not exist") |
| 346 | # ! Send all vulnerabilities to Graylog ! # |
| 347 | for vulnerability in processed_vulnerabilities: |
| 348 | await event_shipper( |
| 349 | EventShipperPayload( |
| 350 | integration="vulnerabilities", |
| 351 | customer_code=customer_code, |
| 352 | agent_name=agent_name, |
| 353 | **vulnerability.model_dump(), |
| 354 | ), |
| 355 | ) |
| 356 | return True |