| 1 | import time |
| 2 | from typing import Any |
| 3 | from typing import Dict |
| 4 | |
| 5 | from cortex4py.api import Api |
| 6 | from fastapi import HTTPException |
| 7 | from loguru import logger |
| 8 | |
| 9 | from app.connectors.cortex.schema.analyzers import AnalyzerJobData |
| 10 | from app.connectors.utils import get_connector_info_from_db |
| 11 | from app.db.db_session import get_db_session |
| 12 | |
| 13 | |
| 14 | async def verify_cortex_credentials(attributes: Dict[str, Any]) -> Dict[str, Any]: |
| 15 | """ |
| 16 | Verifies the connection to Cortex service. |
| 17 | |
| 18 | Returns: |
| 19 | dict: A dictionary containing 'connectionSuccessful' status and 'authToken' if the connection is successful. |
| 20 | """ |
| 21 | logger.info(f"Verifying the Cortex connection to {attributes['connector_url']}") |
| 22 | |
| 23 | try: |
| 24 | api = Api( |
| 25 | attributes["connector_url"], |
| 26 | attributes["connector_api_key"], |
| 27 | verify_cert=False, |
| 28 | ) |
| 29 | # Get Cortex Status |
| 30 | status = api.status |
| 31 | if status: |
| 32 | logger.debug("Cortex connection successful") |
| 33 | return { |
| 34 | "connectionSuccessful": True, |
| 35 | "message": "Cortex connection successful", |
| 36 | } |
| 37 | else: |
| 38 | logger.error( |
| 39 | f"Connection to {attributes['connector_url']} failed with error.", |
| 40 | ) |
| 41 | return { |
| 42 | "connectionSuccessful": False, |
| 43 | "message": f"Connection to {attributes['connector_url']} failed with error.", |
| 44 | } |
| 45 | except Exception as e: |
| 46 | logger.error( |
| 47 | f"Connection to {attributes['connector_url']} failed with error: {e}", |
| 48 | ) |
| 49 | return { |
| 50 | "connectionSuccessful": False, |
| 51 | "message": f"Connection to {attributes['connector_url']} failed with error: {e}", |
| 52 | } |
| 53 | |
| 54 | |
| 55 | async def verify_cortex_connection(connector_name: str) -> str: |
| 56 | """ |
| 57 | Returns the authentication token for the Cortex service. |
| 58 | |
| 59 | Returns: |
| 60 | str: Authentication token for the Cortex service. |
| 61 | """ |
| 62 | async with get_db_session() as session: # This will correctly enter the context manager |
| 63 | attributes = await get_connector_info_from_db(connector_name, session) |
| 64 | if attributes is None: |
| 65 | logger.error("No Cortex connector found in the database") |
| 66 | return None |
| 67 | return await verify_cortex_credentials(attributes) |
| 68 | |
| 69 | |
| 70 | async def create_cortex_client(connector_name: str) -> Api: |
| 71 | """ |
| 72 | Returns an Cortex client for the Wazuh Indexer service. |
| 73 | |
| 74 | Returns: |
| 75 | Cortex: Cortex client for the Cortex service. |
| 76 | """ |
| 77 | async with get_db_session() as session: # This will correctly enter the context manager |
| 78 | attributes = await get_connector_info_from_db(connector_name, session) |
| 79 | if attributes is None: |
| 80 | logger.error("No Wazuh Indexer connector found in the database") |
| 81 | return None |
| 82 | return Api( |
| 83 | attributes["connector_url"], |
| 84 | attributes["connector_api_key"], |
| 85 | verify_cert=False, |
| 86 | ) |
| 87 | |
| 88 | |
| 89 | async def run_and_wait_for_analyzer( |
| 90 | analyzer_name: str, |
| 91 | job_data: AnalyzerJobData, |
| 92 | ) -> Dict[str, Any]: |
| 93 | """ |
| 94 | Runs an analyzer by name and waits for the job to complete. |
| 95 | |
| 96 | Args: |
| 97 | analyzer_name (str): The name of the analyzer to run. |
| 98 | job_data (AnalyzerJobData): The data for the analyzer job. |
| 99 | |
| 100 | Returns: |
| 101 | Dict[str, Any]: A dictionary containing the result of the analyzer job. |
| 102 | """ |
| 103 | api = await create_cortex_client("Cortex") # Create Api object |
| 104 | if api is None: |
| 105 | return {"success": False, "message": "API initialization failed"} |
| 106 | try: |
| 107 | job = api.analyzers.run_by_name(analyzer_name, job_data.model_dump(), force=1) |
| 108 | return await monitor_analyzer_job(api, job) |
| 109 | except Exception as e: |
| 110 | raise HTTPException( |
| 111 | status_code=500, |
| 112 | detail=f"Error running analyzer {analyzer_name}: {e}", |
| 113 | ) |
| 114 | |
| 115 | |
| 116 | async def monitor_analyzer_job(api: Api, job: Any) -> Dict[str, Any]: |
| 117 | """ |
| 118 | Monitors the status of an analyzer job and retrieves the final report when the job is completed. |
| 119 | |
| 120 | Args: |
| 121 | api (Api): The API object used to make requests to the Cortex API. |
| 122 | job (Any): The job object representing the analyzer job. |
| 123 | |
| 124 | Returns: |
| 125 | Dict[str, Any]: A dictionary containing the success status and message of the job. |
| 126 | """ |
| 127 | r_json = job.json() |
| 128 | job_id = r_json["id"] |
| 129 | logger.info(f"Job ID is: {job_id}") |
| 130 | |
| 131 | job_state = r_json["status"] |
| 132 | timer = 0 |
| 133 | |
| 134 | while job_state != "Success": |
| 135 | if timer == 60: |
| 136 | logger.error("Job failed to complete after 5 minutes.") |
| 137 | return {"success": False, "message": "Job timed out"} |
| 138 | |
| 139 | timer += 1 |
| 140 | logger.info(f"Timer is: {timer}") |
| 141 | |
| 142 | if job_state == "Failure": |
| 143 | error_message = r_json["errorMessage"] |
| 144 | logger.error(f"Cortex Failure: {error_message}") |
| 145 | return {"success": False, "message": f"Analyzer failed: {error_message}"} |
| 146 | |
| 147 | time.sleep(5) |
| 148 | followup_request = api.jobs.get_by_id(job_id) |
| 149 | r_json = followup_request.json() |
| 150 | job_state = r_json["status"] |
| 151 | |
| 152 | return await retrieve_final_report(api, job_id) |
| 153 | |
| 154 | |
| 155 | async def retrieve_final_report(api: Api, job_id: str) -> Dict[str, Any]: |
| 156 | """ |
| 157 | Retrieves the final report for a given job ID from the Cortex API. |
| 158 | |
| 159 | Args: |
| 160 | api (Api): The Cortex API instance. |
| 161 | job_id (str): The ID of the job. |
| 162 | |
| 163 | Returns: |
| 164 | Dict[str, Any]: A dictionary containing the success status, message, and final report. |
| 165 | """ |
| 166 | report = api.jobs.get_report(job_id).report |
| 167 | final_report = report["full"] |
| 168 | return { |
| 169 | "success": True, |
| 170 | "message": "Analyzer ran successfully", |
| 171 | "report": final_report, |
| 172 | } |