| 1 | from typing import Any |
| 2 | from typing import Dict |
| 3 | from typing import List |
| 4 | from typing import Optional |
| 5 | from typing import Tuple |
| 6 | from typing import Union |
| 7 | |
| 8 | import httpx |
| 9 | |
| 10 | # import pcre2 |
| 11 | import xmltodict |
| 12 | from fastapi import HTTPException |
| 13 | from loguru import logger |
| 14 | |
| 15 | from app.connectors.wazuh_manager.schema.rules import RuleDisable |
| 16 | from app.connectors.wazuh_manager.schema.rules import RuleDisableResponse |
| 17 | from app.connectors.wazuh_manager.schema.rules import RuleEnable |
| 18 | from app.connectors.wazuh_manager.schema.rules import RuleEnableResponse |
| 19 | from app.connectors.wazuh_manager.schema.rules import RuleExcludeRequest |
| 20 | from app.connectors.wazuh_manager.schema.rules import RuleExcludeResponse |
| 21 | from app.connectors.wazuh_manager.schema.rules import WazuhRuleFileContentResponse |
| 22 | from app.connectors.wazuh_manager.schema.rules import WazuhRuleFilesResponse |
| 23 | from app.connectors.wazuh_manager.schema.rules import WazuhRuleFileUploadResponse |
| 24 | from app.connectors.wazuh_manager.schema.rules import WazuhRulesResponse |
| 25 | from app.connectors.wazuh_manager.utils.universal import restart_service |
| 26 | from app.connectors.wazuh_manager.utils.universal import send_get_request |
| 27 | from app.connectors.wazuh_manager.utils.universal import send_put_request |
| 28 | |
| 29 | |
| 30 | async def get_wazuh_rules(**params) -> WazuhRulesResponse: |
| 31 | """ |
| 32 | Fetch Wazuh rules from the Wazuh Manager API. |
| 33 | |
| 34 | Args: |
| 35 | **params: All query parameters passed directly to the API |
| 36 | |
| 37 | Returns: |
| 38 | WazuhRulesResponse: Structured response with rules data |
| 39 | """ |
| 40 | # Filter out None values |
| 41 | clean_params = {k: v for k, v in params.items() if v is not None} |
| 42 | |
| 43 | # Handle list parameters |
| 44 | if "rule_ids" in clean_params and isinstance(clean_params["rule_ids"], list): |
| 45 | clean_params["rule_ids"] = ",".join(map(str, clean_params["rule_ids"])) |
| 46 | if "select" in clean_params and isinstance(clean_params["select"], list): |
| 47 | clean_params["select"] = ",".join(clean_params["select"]) |
| 48 | if "filename" in clean_params and isinstance(clean_params["filename"], list): |
| 49 | clean_params["filename"] = ",".join(clean_params["filename"]) |
| 50 | |
| 51 | try: |
| 52 | response = await send_get_request(endpoint="/rules", params=clean_params) |
| 53 | |
| 54 | if not response.get("success"): |
| 55 | raise HTTPException(status_code=500, detail="Failed to fetch rules from Wazuh API") |
| 56 | |
| 57 | # Extract data from nested response structure |
| 58 | wazuh_data = response.get("data", {}).get("data", {}) |
| 59 | rules = wazuh_data.get("affected_items", []) |
| 60 | total_items = wazuh_data.get("total_affected_items", len(rules)) |
| 61 | |
| 62 | logger.info(f"Retrieved {len(rules)} of {total_items} Wazuh rules") |
| 63 | |
| 64 | return WazuhRulesResponse( |
| 65 | success=True, |
| 66 | message=f"Successfully retrieved {len(rules)} rules", |
| 67 | results=rules, |
| 68 | total_items=total_items, |
| 69 | ) |
| 70 | |
| 71 | except Exception as e: |
| 72 | logger.error(f"Error fetching Wazuh rules: {e}") |
| 73 | raise HTTPException(status_code=500, detail=f"Error fetching rules: {str(e)}") |
| 74 | |
| 75 | |
| 76 | async def get_wazuh_rule_files(**params) -> WazuhRuleFilesResponse: |
| 77 | """ |
| 78 | Fetch Wazuh rule files from the Wazuh Manager API. |
| 79 | |
| 80 | Args: |
| 81 | **params: All query parameters passed directly to the API |
| 82 | |
| 83 | Returns: |
| 84 | WazuhRuleFilesResponse: Structured response with rule files data |
| 85 | |
| 86 | Raises: |
| 87 | HTTPException: If there's an error fetching the rule files |
| 88 | """ |
| 89 | # Filter out None values and prepare parameters |
| 90 | clean_params = {k: v for k, v in params.items() if v is not None} |
| 91 | |
| 92 | # Handle list parameters that need to be joined as comma-separated strings |
| 93 | if "filename" in clean_params and isinstance(clean_params["filename"], list): |
| 94 | clean_params["filename"] = ",".join(clean_params["filename"]) |
| 95 | if "select" in clean_params and isinstance(clean_params["select"], list): |
| 96 | clean_params["select"] = ",".join(clean_params["select"]) |
| 97 | |
| 98 | logger.debug(f"Requesting Wazuh rule files with params: {clean_params}") |
| 99 | |
| 100 | try: |
| 101 | response = await send_get_request(endpoint="/rules/files", params=clean_params) |
| 102 | |
| 103 | # Check if the API request was successful |
| 104 | if not response.get("success"): |
| 105 | error_detail = response.get("message", "Failed to fetch rule files from Wazuh API") |
| 106 | logger.error(f"Wazuh API error: {error_detail}") |
| 107 | raise HTTPException(status_code=500, detail=error_detail) |
| 108 | |
| 109 | # Extract data from nested response structure |
| 110 | wazuh_data = response.get("data", {}).get("data", {}) |
| 111 | rule_files = wazuh_data.get("affected_items", []) |
| 112 | total_items = wazuh_data.get("total_affected_items", len(rule_files)) |
| 113 | |
| 114 | logger.info(f"Retrieved {len(rule_files)} of {total_items} Wazuh rule files") |
| 115 | |
| 116 | return WazuhRuleFilesResponse( |
| 117 | success=True, |
| 118 | message=f"Successfully retrieved {len(rule_files)} rule files", |
| 119 | results=rule_files, |
| 120 | total_items=total_items, |
| 121 | ) |
| 122 | |
| 123 | except HTTPException: |
| 124 | # Re-raise HTTP exceptions as-is |
| 125 | raise |
| 126 | except Exception as e: |
| 127 | logger.error(f"Error fetching Wazuh rule files: {e}") |
| 128 | raise HTTPException(status_code=500, detail=f"Error fetching rule files: {str(e)}") |
| 129 | |
| 130 | |
| 131 | async def get_wazuh_rule_file_content(filename: str, **params) -> WazuhRuleFileContentResponse: |
| 132 | """ |
| 133 | Fetch the content of a specific Wazuh rule file from the Wazuh Manager API. |
| 134 | |
| 135 | Args: |
| 136 | filename: The name of the rule file to fetch content for |
| 137 | **params: All query parameters passed directly to the API |
| 138 | |
| 139 | Returns: |
| 140 | WazuhRuleFileContentResponse: Structured response with rule file content |
| 141 | |
| 142 | Raises: |
| 143 | HTTPException: If there's an error fetching the rule file content |
| 144 | """ |
| 145 | # Filter out None values and prepare parameters |
| 146 | clean_params = {k: v for k, v in params.items() if v is not None} |
| 147 | |
| 148 | # Check if raw content is requested |
| 149 | is_raw = clean_params.get("raw", False) |
| 150 | |
| 151 | logger.debug(f"Requesting Wazuh rule file content for '{filename}' with params: {clean_params}") |
| 152 | |
| 153 | try: |
| 154 | # Handle raw response differently |
| 155 | if is_raw: |
| 156 | # For raw requests, we need to use a modified approach |
| 157 | # Use only the raw parameter to trigger the special handling in send_get_request |
| 158 | raw_params = {"raw": True} |
| 159 | response = await send_get_request(endpoint=f"/rules/files/{filename}", params=raw_params) |
| 160 | |
| 161 | # Check if the API request was successful |
| 162 | if not response.get("success"): |
| 163 | error_detail = response.get("message", f"Failed to fetch raw rule file content for {filename}") |
| 164 | logger.error(f"Wazuh API error: {error_detail}") |
| 165 | |
| 166 | # Handle specific errors |
| 167 | if "not found" in error_detail.lower(): |
| 168 | raise HTTPException(status_code=404, detail=f"Rule file '{filename}' not found") |
| 169 | else: |
| 170 | raise HTTPException(status_code=500, detail=error_detail) |
| 171 | |
| 172 | # For raw responses, the content is in response["data"] |
| 173 | content = response.get("data", "") |
| 174 | logger.info(f"Retrieved raw content for rule file '{filename}' ({len(content)} characters)") |
| 175 | |
| 176 | return WazuhRuleFileContentResponse( |
| 177 | success=True, |
| 178 | message=f"Successfully retrieved raw content for rule file '{filename}'", |
| 179 | filename=filename, |
| 180 | content=content, |
| 181 | is_raw=True, |
| 182 | total_items=None, |
| 183 | ) |
| 184 | |
| 185 | # Handle structured response (non-raw) |
| 186 | response = await send_get_request(endpoint=f"/rules/files/{filename}", params=clean_params) |
| 187 | |
| 188 | # Check if the API request was successful |
| 189 | if not response.get("success"): |
| 190 | error_detail = response.get("message", f"Failed to fetch rule file content for {filename}") |
| 191 | logger.error(f"Wazuh API error: {error_detail}") |
| 192 | |
| 193 | # Handle specific errors |
| 194 | if "not found" in error_detail.lower(): |
| 195 | raise HTTPException(status_code=404, detail=f"Rule file '{filename}' not found") |
| 196 | else: |
| 197 | raise HTTPException(status_code=500, detail=error_detail) |
| 198 | |
| 199 | # Handle structured response |
| 200 | wazuh_data = response.get("data", {}).get("data", {}) |
| 201 | affected_items = wazuh_data.get("affected_items", []) |
| 202 | total_items = wazuh_data.get("total_affected_items", len(affected_items)) |
| 203 | |
| 204 | if not affected_items: |
| 205 | raise HTTPException(status_code=404, detail=f"No content found for rule file '{filename}'") |
| 206 | |
| 207 | # Extract the content from the first affected item |
| 208 | content = affected_items[0] if affected_items else {} |
| 209 | |
| 210 | logger.info(f"Retrieved structured content for rule file '{filename}' with {total_items} affected items") |
| 211 | |
| 212 | return WazuhRuleFileContentResponse( |
| 213 | success=True, |
| 214 | message=f"Successfully retrieved content for rule file '{filename}'", |
| 215 | filename=filename, |
| 216 | content=content, |
| 217 | is_raw=False, |
| 218 | total_items=total_items, |
| 219 | ) |
| 220 | |
| 221 | except HTTPException: |
| 222 | # Re-raise HTTP exceptions as-is |
| 223 | raise |
| 224 | except Exception as e: |
| 225 | logger.error(f"Error fetching Wazuh rule file content for '{filename}': {e}") |
| 226 | raise HTTPException(status_code=500, detail=f"Error fetching rule file content: {str(e)}") |
| 227 | |
| 228 | |
| 229 | async def update_wazuh_rule_file( |
| 230 | filename: str, |
| 231 | file_content: bytes, |
| 232 | pretty: Optional[bool] = False, |
| 233 | wait_for_complete: Optional[bool] = False, |
| 234 | overwrite: Optional[bool] = False, |
| 235 | relative_dirname: Optional[str] = None, |
| 236 | ) -> WazuhRuleFileUploadResponse: |
| 237 | """ |
| 238 | Upload or update a Wazuh rule file. |
| 239 | |
| 240 | Args: |
| 241 | filename: Name of the rule file |
| 242 | file_content: Binary content of the rule file |
| 243 | pretty: Show results in human-readable format |
| 244 | wait_for_complete: Disable timeout response |
| 245 | overwrite: Whether to overwrite the file if it exists |
| 246 | relative_dirname: Relative directory name |
| 247 | |
| 248 | Returns: |
| 249 | WazuhRuleFileUploadResponse: Response indicating success/failure |
| 250 | |
| 251 | Raises: |
| 252 | HTTPException: If there's an error uploading the file |
| 253 | """ |
| 254 | # Prepare parameters |
| 255 | params = {} |
| 256 | if pretty is not None: |
| 257 | params["pretty"] = str(pretty).lower() |
| 258 | if wait_for_complete is not None: |
| 259 | params["wait_for_complete"] = str(wait_for_complete).lower() |
| 260 | if overwrite is not None: |
| 261 | params["overwrite"] = str(overwrite).lower() |
| 262 | if relative_dirname is not None: |
| 263 | params["relative_dirname"] = relative_dirname |
| 264 | |
| 265 | logger.info(f"Uploading/updating rule file: {filename}") |
| 266 | logger.debug(f"Request params: {params}") |
| 267 | |
| 268 | try: |
| 269 | # Send PUT request with binary data |
| 270 | response = await send_put_request( |
| 271 | endpoint=f"/rules/files/{filename}", |
| 272 | data=file_content, |
| 273 | params=params, |
| 274 | binary_data=True, |
| 275 | debug=True, |
| 276 | ) |
| 277 | |
| 278 | # Check if the API request was successful |
| 279 | if not response.get("success"): |
| 280 | error_detail = response.get("message", "Failed to upload rule file to Wazuh API") |
| 281 | status_code = response.get("status_code", 500) |
| 282 | logger.error(f"Wazuh API error: {error_detail}") |
| 283 | raise HTTPException(status_code=status_code, detail=error_detail) |
| 284 | |
| 285 | # Extract data from response |
| 286 | wazuh_data = response.get("data", {}).get("data", {}) |
| 287 | total_items = wazuh_data.get("total_affected_items", 1) |
| 288 | |
| 289 | logger.info(f"Successfully uploaded/updated rule file: {filename}") |
| 290 | |
| 291 | return WazuhRuleFileUploadResponse( |
| 292 | success=True, |
| 293 | message=f"Successfully uploaded/updated rule file: {filename}", |
| 294 | filename=filename, |
| 295 | details=wazuh_data, |
| 296 | total_items=total_items, |
| 297 | ) |
| 298 | |
| 299 | except HTTPException: |
| 300 | # Re-raise HTTP exceptions as-is |
| 301 | raise |
| 302 | except Exception as e: |
| 303 | logger.error(f"Error uploading rule file {filename}: {e}") |
| 304 | raise HTTPException(status_code=500, detail=f"Error uploading rule file: {str(e)}") |
| 305 | |
| 306 | |
| 307 | async def fetch_filename(rule_id: str) -> str: |
| 308 | """ |
| 309 | Fetches the filename associated with a given rule ID from the Wazuh Manager. |
| 310 | |
| 311 | Args: |
| 312 | rule_id (str): The ID of the rule. |
| 313 | |
| 314 | Returns: |
| 315 | str: The filename associated with the rule ID. |
| 316 | |
| 317 | Raises: |
| 318 | HTTPException: If the rule ID is not found in the Wazuh Manager. |
| 319 | """ |
| 320 | endpoint = "rules" |
| 321 | params = {"rule_ids": rule_id} |
| 322 | filename_data = await send_get_request(endpoint=endpoint, params=params) |
| 323 | if filename_data["data"]["data"]["total_affected_items"] == 0: |
| 324 | raise HTTPException( |
| 325 | status_code=404, |
| 326 | detail=f"Rule {rule_id} not found. Make sure the rule ID is correct within the Wazuh Manager.", |
| 327 | ) |
| 328 | return filename_data["data"]["data"]["affected_items"][0]["filename"] |
| 329 | |
| 330 | |
| 331 | async def fetch_file_content(filename: str) -> str: |
| 332 | """ |
| 333 | Fetches the content of a file from the Wazuh Manager. |
| 334 | |
| 335 | Args: |
| 336 | filename (str): The name of the file to fetch. |
| 337 | |
| 338 | Returns: |
| 339 | str: The content of the file. |
| 340 | |
| 341 | Raises: |
| 342 | HTTPException: If the file is not found in the Wazuh Manager. |
| 343 | """ |
| 344 | endpoint = f"rules/files/{filename}" |
| 345 | file_content_data = await send_get_request(endpoint=endpoint) |
| 346 | if file_content_data["data"]["data"]["total_affected_items"] == 0: |
| 347 | raise HTTPException( |
| 348 | status_code=404, |
| 349 | detail=f"File {filename} not found. Make sure the file name is correct within the Wazuh Manager.", |
| 350 | ) |
| 351 | return file_content_data["data"]["data"]["affected_items"][0]["group"] |
| 352 | |
| 353 | |
| 354 | async def set_rule_level( |
| 355 | file_content: Any, |
| 356 | rule_id: str, |
| 357 | new_level: str, |
| 358 | ) -> Tuple[str, Any]: |
| 359 | """ |
| 360 | Sets the level of a rule identified by its ID in the given file content. |
| 361 | |
| 362 | Args: |
| 363 | file_content (Any): The content of the file containing the rules. |
| 364 | rule_id (str): The ID of the rule to set the level for. |
| 365 | new_level (str): The new level to set for the rule. |
| 366 | |
| 367 | Returns: |
| 368 | Tuple[str, Any]: A tuple containing the previous level of the rule and the modified file content. |
| 369 | """ |
| 370 | previous_level = None |
| 371 | try: |
| 372 | if isinstance(file_content, dict): |
| 373 | file_content = [file_content] |
| 374 | for group_block in file_content: |
| 375 | rule_block = group_block.get("rule", None) |
| 376 | if rule_block: |
| 377 | if isinstance(rule_block, dict): |
| 378 | rule_block = [rule_block] |
| 379 | for rule in rule_block: |
| 380 | if rule["@id"] == rule_id: |
| 381 | previous_level = rule["@level"] |
| 382 | rule["@level"] = new_level |
| 383 | break |
| 384 | except (KeyError, TypeError) as e: |
| 385 | raise HTTPException(status_code=500, detail=f"Failed to set rule level: {e}") |
| 386 | return previous_level, file_content |
| 387 | |
| 388 | |
| 389 | async def convert_to_xml( |
| 390 | updated_file_content: Union[Dict[str, str], List[Dict[str, str]]], |
| 391 | ) -> str: |
| 392 | """ |
| 393 | Converts the updated file content to XML format. |
| 394 | |
| 395 | Args: |
| 396 | updated_file_content (Union[Dict[str, str], List[Dict[str, str]]]): The updated file content. |
| 397 | |
| 398 | Returns: |
| 399 | str: The XML content. |
| 400 | |
| 401 | Raises: |
| 402 | HTTPException: If there is an error converting to XML. |
| 403 | """ |
| 404 | xml_content_list = [] |
| 405 | try: |
| 406 | for group in updated_file_content: |
| 407 | xml_dict = {"group": group} |
| 408 | xml_content = xmltodict.unparse(xml_dict, pretty=True) |
| 409 | xml_content = xml_content.replace( |
| 410 | '<?xml version="1.0" encoding="utf-8"?>', |
| 411 | "", |
| 412 | ) |
| 413 | xml_content_list.append(xml_content) |
| 414 | except Exception as e: |
| 415 | raise HTTPException(status_code=500, detail=f"Failed to convert to XML: {e}") |
| 416 | xml_content = "\n".join(xml_content_list) |
| 417 | xml_content = xml_content.strip() |
| 418 | return xml_content |
| 419 | |
| 420 | |
| 421 | async def upload_updated_rule(filename: str, xml_content: str): |
| 422 | """ |
| 423 | Uploads an updated rule to the Wazuh Manager. |
| 424 | |
| 425 | Args: |
| 426 | filename (str): The name of the rule file. |
| 427 | xml_content (str): The content of the rule file in XML format. |
| 428 | |
| 429 | Returns: |
| 430 | dict: The response from the Wazuh Manager API. |
| 431 | |
| 432 | Raises: |
| 433 | HTTPException: If the upload fails. |
| 434 | """ |
| 435 | response = await send_put_request( |
| 436 | endpoint=f"rules/files/{filename}", |
| 437 | data=xml_content, |
| 438 | params={"overwrite": "true"}, |
| 439 | ) |
| 440 | logger.info(response) |
| 441 | if response["data"]["data"]["total_affected_items"] == 0: |
| 442 | raise HTTPException( |
| 443 | status_code=500, |
| 444 | detail="Failed to upload updated rule to Wazuh Manager.", |
| 445 | ) |
| 446 | return response |
| 447 | |
| 448 | |
| 449 | async def process_rule(rule, rule_action_func, ResponseModel): |
| 450 | """ |
| 451 | Process a rule by fetching its filename and content, applying a rule action function, |
| 452 | converting the updated content to XML, uploading the updated rule, and restarting the service. |
| 453 | |
| 454 | Args: |
| 455 | rule: The rule to be processed. |
| 456 | rule_action_func: The function to apply to the rule's content. |
| 457 | ResponseModel: The response model class. |
| 458 | |
| 459 | Returns: |
| 460 | An instance of ResponseModel with the previous level, success status, and a message. |
| 461 | """ |
| 462 | filename, file_content = await fetch_filename_and_content(rule.rule_id) |
| 463 | previous_level, updated_file_content = await rule_action_func( |
| 464 | file_content, |
| 465 | rule.rule_id, |
| 466 | ) |
| 467 | xml_content = await convert_to_xml(updated_file_content) |
| 468 | await upload_updated_rule(filename, xml_content) |
| 469 | await restart_service() |
| 470 | return ResponseModel( |
| 471 | previous_level=previous_level, |
| 472 | success=True, |
| 473 | message=f"Rule {rule.rule_id} successfully processed in file {filename}.", |
| 474 | ) |
| 475 | |
| 476 | |
| 477 | async def fetch_filename_and_content(rule_id: str) -> Tuple[str, str]: |
| 478 | """ |
| 479 | Fetches the filename and content of a rule based on the given rule ID. |
| 480 | |
| 481 | Args: |
| 482 | rule_id (str): The ID of the rule. |
| 483 | |
| 484 | Returns: |
| 485 | Tuple[str, str]: A tuple containing the filename and content of the rule. |
| 486 | """ |
| 487 | filename = await fetch_filename(rule_id) |
| 488 | file_content = await fetch_file_content(filename) |
| 489 | return filename, file_content |
| 490 | |
| 491 | |
| 492 | async def disable_rule(rule: RuleDisable) -> RuleDisableResponse: |
| 493 | """ |
| 494 | Disable a rule by setting its level to "1". |
| 495 | |
| 496 | Args: |
| 497 | rule (RuleDisable): The rule to be disabled. |
| 498 | |
| 499 | Returns: |
| 500 | RuleDisableResponse: The response indicating the success or failure of the operation. |
| 501 | """ |
| 502 | |
| 503 | async def process(fc, rid): |
| 504 | return await set_rule_level(fc, rid, "1") |
| 505 | |
| 506 | return await process_rule(rule, process, RuleDisableResponse) |
| 507 | |
| 508 | |
| 509 | async def enable_rule(rule: RuleEnable, previous_level: str) -> RuleEnableResponse: |
| 510 | """ |
| 511 | Enable a rule with the given parameters. |
| 512 | |
| 513 | Args: |
| 514 | rule (RuleEnable): The rule to enable. |
| 515 | previous_level (str): The previous level of the rule. |
| 516 | |
| 517 | Returns: |
| 518 | RuleEnableResponse: The response indicating the success or failure of enabling the rule. |
| 519 | """ |
| 520 | |
| 521 | async def process(fc, rid): |
| 522 | return await set_rule_level(fc, rid, previous_level) |
| 523 | |
| 524 | return await process_rule(rule, process, RuleEnableResponse) |
| 525 | |
| 526 | |
| 527 | ################# ! EXCLUDE RULE ! ################# |
| 528 | async def post_to_copilot_ai_module(data: RuleExcludeRequest) -> RuleExcludeResponse: |
| 529 | """ |
| 530 | Send a POST request to the copilot-ai-module Docker container. |
| 531 | |
| 532 | Args: |
| 533 | data (CollectHuntress): The data to send to the copilot-ai-module Docker container. |
| 534 | """ |
| 535 | logger.info(f"Sending POST request to http://copilot-ai-module/wazuh-rule-exclusion with data: {data.model_dump()}") |
| 536 | # raise HTTPException(status_code=501, detail="Not Implemented Yet") |
| 537 | async with httpx.AsyncClient() as client: |
| 538 | data = await client.post( |
| 539 | "http://copilot-ai-module/wazuh-rule-exclusion", |
| 540 | json=data.model_dump(), |
| 541 | timeout=120, |
| 542 | ) |
| 543 | return RuleExcludeResponse(**data.json()) |