|
| 1 | +import demistomock as demisto # noqa: F401 |
| 2 | +from CommonServerPython import * # noqa: F401 |
| 3 | + |
| 4 | +""" |
| 5 | +Cortex XSOAR Integration for Azure OpenAI |
| 6 | +""" |
| 7 | + |
| 8 | +import traceback |
| 9 | +import json |
| 10 | + |
| 11 | + |
| 12 | +# --- CONSTANTS --- |
| 13 | +API_VERSION = "2024-02-01" # Using a stable, recent API version |
| 14 | +DEFAULT_DEPLOYMENT = "gpt-4o" |
| 15 | + |
| 16 | + |
| 17 | +class Client(BaseClient): |
| 18 | + """ |
| 19 | + Client class to interact with the Azure OpenAI service. |
| 20 | + Inherits from BaseClient from CommonServerPython to handle SSL verification, etc. |
| 21 | + """ |
| 22 | + |
| 23 | + def __init__(self, server_url, api_key, deployment_name, instruction, verify): |
| 24 | + # Ensure the URL ends with a / |
| 25 | + base_url = server_url if server_url.endswith("/") else f"{server_url}/" |
| 26 | + super().__init__(base_url=base_url, verify=verify) |
| 27 | + self.api_key = api_key |
| 28 | + self.deployment_name = deployment_name |
| 29 | + self.instruction = instruction |
| 30 | + |
| 31 | + # FIX: Added 'require_json' parameter to conditionally request JSON format. |
| 32 | + def send_message(self, message: str, require_json: bool = False) -> dict: |
| 33 | + """ |
| 34 | + Sends a message to the Azure OpenAI 'chat completions' endpoint. |
| 35 | +
|
| 36 | + :param message: The user message to send. |
| 37 | + :param require_json: If True, requests the response in JSON object format. |
| 38 | + :return: The full JSON response from the API. |
| 39 | + """ |
| 40 | + full_url_path = f"openai/deployments/{self.deployment_name}/chat/completions?api-version={API_VERSION}" |
| 41 | + headers = {"Content-Type": "application/json", "api-key": self.api_key} |
| 42 | + payload = { |
| 43 | + "messages": [{"role": "system", "content": self.instruction}, {"role": "user", "content": message}], |
| 44 | + "max_tokens": 2000, |
| 45 | + "temperature": 0.5, |
| 46 | + "top_p": 0.95, |
| 47 | + "frequency_penalty": 0, |
| 48 | + "presence_penalty": 0, |
| 49 | + "stop": None, |
| 50 | + "stream": False, |
| 51 | + } |
| 52 | + |
| 53 | + # FIX: Only add the 'response_format' parameter when it's explicitly required. |
| 54 | + if require_json: |
| 55 | + payload["response_format"] = {"type": "json_object"} |
| 56 | + |
| 57 | + # The _http_request method handles the HTTP request |
| 58 | + response = self._http_request( |
| 59 | + method="POST", url_suffix=full_url_path, headers=headers, json_data=payload, resp_type="json" |
| 60 | + ) |
| 61 | + return response |
| 62 | + |
| 63 | + |
| 64 | +# --- COMMAND FUNCTIONS --- |
| 65 | + |
| 66 | + |
| 67 | +def test_module(client: Client) -> str: |
| 68 | + """ |
| 69 | + Tests API connectivity by sending a simple message. |
| 70 | + """ |
| 71 | + try: |
| 72 | + # Use a simple instruction for the test. |
| 73 | + client.instruction = "You are a helpful assistant. Respond with 'ok' if you are working." |
| 74 | + # FIX: Call send_message without requiring a JSON response. |
| 75 | + client.send_message("This is a connection test.", require_json=False) |
| 76 | + return "ok" |
| 77 | + except DemistoException as e: |
| 78 | + if "401" in str(e): |
| 79 | + return "Authentication Error: Check your API Key." |
| 80 | + elif "404" in str(e): |
| 81 | + return "Connection Error: Check the Endpoint URL and Deployment Name." |
| 82 | + else: |
| 83 | + # Provide the specific error message from the API for easier debugging. |
| 84 | + return f"An API error occurred: {str(e)}" |
| 85 | + |
| 86 | + |
| 87 | +def send_message_command(client: Client, args: dict) -> CommandResults: |
| 88 | + """ |
| 89 | + Executes the command to send a message and process the structured response. |
| 90 | + """ |
| 91 | + message = args.get("message", "") |
| 92 | + |
| 93 | + demisto.debug(f"Sending message to Azure OpenAI: '{message}'") |
| 94 | + # FIX: Call send_message *with* the requirement for a JSON response. |
| 95 | + raw_response = client.send_message(message, require_json=True) |
| 96 | + |
| 97 | + # Extract the response content |
| 98 | + raw_answer = "" |
| 99 | + if raw_response.get("choices") and isinstance(raw_response["choices"], list) and len(raw_response["choices"]) > 0: |
| 100 | + first_choice = raw_response["choices"][0] |
| 101 | + if first_choice.get("message") and first_choice.get("message").get("content"): |
| 102 | + raw_answer = first_choice["message"]["content"] |
| 103 | + |
| 104 | + if not raw_answer: |
| 105 | + raise DemistoException("Could not get a valid response from the API.", res=raw_response) |
| 106 | + |
| 107 | + # Logic to handle the structured JSON response |
| 108 | + try: |
| 109 | + # Attempt to parse the response as JSON |
| 110 | + parsed_data = json.loads(raw_answer) |
| 111 | + |
| 112 | + # Map the parsed data to the desired output keys |
| 113 | + outputs = { |
| 114 | + "IncidentAIVerdict": parsed_data.get("IncidentAIVerdict"), |
| 115 | + "AISummary": parsed_data.get("AISummary"), |
| 116 | + "Justification": parsed_data.get("Justification"), |
| 117 | + "ConfidenceScore": parsed_data.get("ConfidenceScore"), |
| 118 | + "EmailHeaderAIAnalysis": parsed_data.get("EmailHeaderAIAnalysis"), |
| 119 | + "EmailAISummary": parsed_data.get("EmailAISummary"), |
| 120 | + "EmailAIVerdict": parsed_data.get("EmailAIVerdict"), |
| 121 | + "Answer": parsed_data, |
| 122 | + } |
| 123 | + |
| 124 | + # Create a formatted, human-readable output |
| 125 | + readable_output = "## 🤖 AI Analysis\n\n" |
| 126 | + readable_output += f"**Verdict:** {outputs.get('IncidentAIVerdict', 'N/A')}\n" |
| 127 | + readable_output += f"**Confidence:** {outputs.get('ConfidenceScore', 'N/A')}\n" |
| 128 | + readable_output += f"**Summary:**\n---\n{outputs.get('AISummary', 'N/A')}\n\n" |
| 129 | + readable_output += f"**Justification:**\n---\n{outputs.get('Justification', 'N/A')}\n\n" |
| 130 | + readable_output += "### Detailed Email Analysis\n" |
| 131 | + readable_output += f"**Email Verdict:** {outputs.get('EmailAIVerdict', 'N/A')}\n" |
| 132 | + readable_output += f"**Email Summary:**\n---\n{outputs.get('EmailAISummary', 'N/A')}\n\n" |
| 133 | + readable_output += f"**Header Analysis:**\n---\n{outputs.get('EmailHeaderAIAnalysis', 'N/A')}\n" |
| 134 | + |
| 135 | + # Update the incident's custom fields |
| 136 | + custom_fields_to_set = { |
| 137 | + "incidentaiverdict": outputs.get("IncidentAIVerdict"), |
| 138 | + "aisummary": outputs.get("AISummary"), |
| 139 | + "justification": outputs.get("Justification"), |
| 140 | + "confidencescore": outputs.get("ConfidenceScore"), |
| 141 | + "emailheaderaianalysis": outputs.get("EmailHeaderAIAnalysis"), |
| 142 | + "emailaisummary": outputs.get("EmailAISummary"), |
| 143 | + "emailaiverdict": outputs.get("EmailAIVerdict"), |
| 144 | + } |
| 145 | + filtered_custom_fields = {k: v for k, v in custom_fields_to_set.items() if v is not None} |
| 146 | + if filtered_custom_fields: |
| 147 | + try: |
| 148 | + demisto.executeCommand("setIncident", {"customFields": filtered_custom_fields}) |
| 149 | + except Exception as e: |
| 150 | + demisto.debug(f"Could not set incident fields. Error: {e}") |
| 151 | + |
| 152 | + except json.JSONDecodeError: |
| 153 | + demisto.debug("AI response was not valid JSON. Treating as raw text.") |
| 154 | + outputs = {"Answer": raw_answer} |
| 155 | + readable_output = f"## Azure OpenAI Response (Raw Text)\n\n**Response:**\n---\n{raw_answer}\n" |
| 156 | + try: |
| 157 | + demisto.executeCommand("setIncident", {"customFields": {"aianswer": raw_answer}}) |
| 158 | + except Exception as e: |
| 159 | + demisto.debug(f"Could not set the 'aianswer' incident field. Error: {e}") |
| 160 | + |
| 161 | + return CommandResults( |
| 162 | + readable_output=readable_output, |
| 163 | + outputs_prefix="AzureOpenAI.Analysis", |
| 164 | + outputs_key_field="IncidentAIVerdict", |
| 165 | + outputs=outputs, |
| 166 | + raw_response=raw_response, |
| 167 | + ) |
| 168 | + |
| 169 | + |
| 170 | +# --- MAIN FUNCTION --- |
| 171 | + |
| 172 | + |
| 173 | +def main() -> None: |
| 174 | + """ |
| 175 | + Main function, parses parameters and executes commands. |
| 176 | + """ |
| 177 | + params = demisto.params() |
| 178 | + server_url = params.get("url") |
| 179 | + |
| 180 | + api_key_details = params.get("credentials").get("password") |
| 181 | + api_key = api_key_details.get("password") if isinstance(api_key_details, dict) else api_key_details |
| 182 | + |
| 183 | + instruction = params.get("instruction") |
| 184 | + if not instruction or "{" not in instruction: |
| 185 | + demisto.debug("Instruction does not seem to request JSON. Using a default prompt for email analysis.") |
| 186 | + instruction = """ |
| 187 | + You are an expert cybersecurity analyst. Analyze the provided data. |
| 188 | + Your response MUST be a single, valid JSON object. Do NOT provide ANY text outside of the JSON object. |
| 189 | + The JSON format must be as follows: |
| 190 | + { |
| 191 | + "IncidentAIVerdict": "string (Malicious, Suspicious, Benign, Informational)", |
| 192 | + "AISummary": "string (A 2-3 sentence global summary of the incident.)", |
| 193 | + "Justification": "string (The primary reason for the verdict, based on the strongest evidence.)", |
| 194 | + "ConfidenceScore": "integer (A confidence score from 0 to 100 for the verdict.)", |
| 195 | + "EmailHeaderAIAnalysis": "string (A detailed analysis of the email headers, including SPF/DKIM/DMARC and routing.)", |
| 196 | + "EmailAISummary": "string (A summary of the analysis of the email body, URLs, and attachments.)", |
| 197 | + "EmailAIVerdict": "string (Phishing, Malware, Spam, BEC, Safe)" |
| 198 | + } |
| 199 | + """ |
| 200 | + |
| 201 | + deployment_name = params.get("deployment_name") or DEFAULT_DEPLOYMENT |
| 202 | + verify_certificate = not params.get("insecure", False) |
| 203 | + |
| 204 | + command = demisto.command() |
| 205 | + args = demisto.args() |
| 206 | + demisto.debug(f"Command being executed is {command}") |
| 207 | + |
| 208 | + try: |
| 209 | + client = Client( |
| 210 | + server_url=server_url, |
| 211 | + api_key=api_key, |
| 212 | + deployment_name=deployment_name, |
| 213 | + instruction=instruction, |
| 214 | + verify=verify_certificate, |
| 215 | + ) |
| 216 | + |
| 217 | + if command == "test-module": |
| 218 | + return_results(test_module(client)) |
| 219 | + elif command == "azure-openai-send-message": |
| 220 | + return_results(send_message_command(client, args)) |
| 221 | + else: |
| 222 | + raise NotImplementedError(f"Command '{command}' is not implemented.") |
| 223 | + |
| 224 | + except Exception as e: |
| 225 | + demisto.error(traceback.format_exc()) |
| 226 | + return_error(f"Failed to execute {command} command. Error: {str(e)}") |
| 227 | + |
| 228 | + |
| 229 | +# --- ENTRY POINT --- |
| 230 | + |
| 231 | +if __name__ in ("__main__", "__builtin__", "builtins"): |
| 232 | + main() |
0 commit comments