import demistomock as demisto # pylint: disable=W9005 from CommonServerPython import * # pylint: disable=W9011 from CommonServerUserPython import * """IMPORTS""" from datetime import date, datetime, timedelta from datetime import time as dttime from decimal import Decimal import snowflake.connector # pylint: disable=E0401,E0611 from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import serialization """GLOBAL VARS""" PARAMS = demisto.params() # pylint: disable=W9016 CREDENTIALS = PARAMS.get("credentials", {}) USER = CREDENTIALS.get("identifier") PASSWORD = CREDENTIALS.get("password") CERTIFICATE = CREDENTIALS.get("credentials", {}).get("sshkey", "").encode() CERT_PASSWORD = CREDENTIALS.get("credentials", {}).get("password") CERT_PASSWORD = CERT_PASSWORD.encode() if CERT_PASSWORD else None ACCOUNT = PARAMS.get("account") AUTHENTICATOR = PARAMS.get("authenticator") REGION = PARAMS.get("region") WAREHOUSE = PARAMS.get("warehouse") DATABASE = PARAMS.get("database") SCHEMA = PARAMS.get("schema") ROLE = PARAMS.get("role") INSECURE = PARAMS.get("insecure", False) OAUTH_CLIENT_ID = PARAMS.get("oauth_client_id") OAUTH_CLIENT_SECRET = PARAMS.get("oauth_client_secret", {}).get("password") OAUTH_TOKEN_URL = PARAMS.get("oauth_token_url") OAUTH_SCOPE_RAW = PARAMS.get("oauth_scope", "") def parse_oauth_scope(raw_scope: str) -> str | None: """Convert a comma-separated scope string to space-separated as required by the Snowflake API. Args: raw_scope: Comma-separated scope string (e.g. "scope1,scope2,scope3"). Returns: Space-separated scope string, or None if the input is empty. """ if not raw_scope: return None result = " ".join(scope.strip() for scope in raw_scope.split(",") if scope.strip()) return result or None OAUTH_SCOPE = parse_oauth_scope(OAUTH_SCOPE_RAW) # How much time before the first fetch to retrieve incidents IS_FETCH = PARAMS.get("isFetch") FETCH_TIME = PARAMS.get("fetch_time") FETCH_QUERY = PARAMS.get("fetch_query") DATETIME_COLUMN = PARAMS.get("datetime_column") INCIDENT_NAME_COLUMN = PARAMS.get("incident_name_column") MAX_ROWS = int(PARAMS.get("limit")) if PARAMS.get("limit") else 10000 TYPE_CODE_TO_DATATYPE = { 0: "number/int", 1: "real", 2: "varchar/string", 3: "date", 4: "timestamp", 5: "variant", 6: "timestamp_ltz", 7: "timestamp_tz", 8: "timestamp_tz", 9: "object", 10: "array", 11: "binary", 12: "time", 13: "boolean", } DT_NEEDS_CHECKING = {"date", "timestamp", "timestamp_ltz", "timestamp_tz", "time"} """SETUP""" if IS_FETCH and not (FETCH_QUERY and DATETIME_COLUMN): err_msg = "When fetching is enabled there are two additional parameters that are required;" err_msg += " The fetch query that determines what data to fetch and the name of the column" err_msg += " in the fetched data that contains a datetime object or timestamp." raise Exception(err_msg) """HELPER FUNCTIONS""" def convert_datetime_to_string(v): # pylint: disable=W9014 """ Parses date, time, timedelta, or datetime object into string parameter: (datetime/date/time/timedelta) v The datetime/date/time/timedelta object to convert returns: Formatted string of the object """ if isinstance(v, datetime): formatted_datetime = v.strftime("%Y-%m-%d %H:%M:%S.%f %z").strip() if "." in formatted_datetime: # convert the time format to 2 digits after the decimal point to co-p with server formats. split_datetime = formatted_datetime.split(".") fixed_datetime = split_datetime[0] + "." + split_datetime[1][:2] else: fixed_datetime = formatted_datetime return fixed_datetime.strip() elif isinstance(v, date): return v.strftime("%Y-%m-%d").strip() elif isinstance(v, dttime): return v.strftime("%H:%M:%S.%f").strip() return v def error_message_from_snowflake_error(e): # pylint: disable=W9014 """ Return formatted error message from contents of a Snowflake error parameter: (snowflake.connector.errors.Error) e The Snowflake error object returns: Formatted error message """ err_msg = f"Snowflake DB error code: {e.errno}\n" err_msg += f"ANSI-compliant SQL State code: {e.sqlstate}\n" err_msg += f"Snowflake query ID: {e.sfqid}\n" err_msg += "Error message: {}" if e.errno == 606: first_sentence = e.raw_msg[: e.raw_msg.find(".") + 1] err_msg = err_msg.format(first_sentence) err_msg += " Specify an active warehouse in the command " err_msg += "arguments or in the integration parameters." elif e.errno == 2003: err_msg = err_msg.format(e.raw_msg) err_msg += " A possible explanation is that the values you entered" err_msg += " for the 'warehouse' and 'database' were incorrect." else: err_msg = err_msg.format(e.raw_msg) return err_msg def set_provided(params, key, val1, val2=None): # pylint: disable=W9014 """ If value is provided, set it in the dict """ if val1: params[key] = val1 elif val2: params[key] = val2 def process_table_row(row, checks): # pylint: disable=W9014 """ Check row data and reformat if necessary The 'checks' parameter contains the names of fields that have the potential to cause issues when they will be json decoded. This function checks the values of the fields flagged in the 'checks' parameter and formats the contents to a json friendly type if necessary. parameter: (dict) row The data (table row) that needs to be processed parameter: (dict[str, list]) checks Dictionary where the key is a string indicative of the type (or bucket of types) that needs reformatting and the values are a list of column names whose data is of that type returns: Reformatted Row """ for column_name, val in row.items(): if column_name in checks.get("isDecimal", []): # Then check the value and reformat it if necessary if isinstance(val, Decimal): row[column_name] = str(val) elif column_name in checks.get("isDT", []): # Then reformat it if necessary row[column_name] = convert_datetime_to_string(val) return row def format_to_json_serializable(column_descriptions, results): # pylint: disable=W9014 """ Screen and reformat any data in 'results' argument that is not json serializable, and return 'results'. 'results' can be a table of data (a list of rows) or a single row. parameter: (list) column_descriptions The metadata that describes data for each column in the 'results' parameter parameter: (list/dict) results What was returned by the cursor object's execute or fetch operation returns: Reformatted 'results' """ name = 0 type_code = 1 checks: dict = {} # Screen by type_code for col in column_descriptions: # if col[type_code] == 0: if TYPE_CODE_TO_DATATYPE.get(col[type_code]) == "number/int": # Then need to check that column's data to see if its data type is Decimal checks.setdefault("isDecimal", []).append(col[name]) # elif col[type_code] in {3, 4, 6, 7, 8, 12}: elif TYPE_CODE_TO_DATATYPE.get(col[type_code]) in DT_NEEDS_CHECKING: # Then need to check that column's data to see if its data type is date, time, timedelta or datetime checks.setdefault("isDT", []).append(col[name]) # if 'results' is a list then it is a data table (list of rows) and need to process each row # in the table, otherwise if 'results' is a dict then it a single table row # Check candidates and reformat if necessary if isinstance(results, dict): results = process_table_row(results, checks) else: # if 'results' isn't a dict, assume it's a list for i, row in enumerate(results): results[i] = process_table_row(row, checks) return results def get_connection_params(args): # pylint: disable=W9014 """ Construct and return the connection parameters parameter: (dict) args The command arguments of the command function calling this helper function returns: Snowflake connection params """ params: dict = {} # Validate that at least one authentication method is configured has_oauth = any([OAUTH_CLIENT_ID, OAUTH_TOKEN_URL, OAUTH_CLIENT_SECRET, OAUTH_SCOPE]) has_certificate = bool(CERTIFICATE) has_password = bool(PASSWORD) if not any([has_oauth, has_certificate, has_password]): raise ValueError( "No authentication method configured. Please provide one of: " "Credentials (username/password), Certificate (key pair), or OAuth parameters." ) if not USER: raise ValueError("Username is required for all authentication methods (username/password, key pair, and OAuth).") set_provided(params, "user", USER) # user value is required for all authentication methods set_provided(params, "account", ACCOUNT) set_provided(params, "region", REGION) set_provided(params, "insecure_mode", INSECURE) set_provided(params, "warehouse", args.get("warehouse"), WAREHOUSE) set_provided(params, "database", args.get("database"), DATABASE) set_provided(params, "schema", args.get("schema"), SCHEMA) set_provided(params, "role", args.get("role"), ROLE) if has_oauth: if not all([OAUTH_CLIENT_ID, OAUTH_TOKEN_URL, OAUTH_CLIENT_SECRET]): raise ValueError("OAuth Client ID, Client Secret, and Token URL must all be provided for OAuth authentication.") # ── OAuth (External OAuth, Client Credentials grant) ── # The snowflake-connector-python natively supports the client credentials flow. # Override the authenticator with the OAuth-specific value. set_provided(params, "authenticator", "OAUTH_CLIENT_CREDENTIALS") set_provided(params, "oauth_client_id", OAUTH_CLIENT_ID) set_provided(params, "oauth_client_secret", OAUTH_CLIENT_SECRET) set_provided(params, "oauth_token_request_url", OAUTH_TOKEN_URL) set_provided(params, "oauth_scope", OAUTH_SCOPE) elif has_certificate: set_provided(params, "authenticator", AUTHENTICATOR) # ── Key Pair authentication ── p_key = serialization.load_pem_private_key(CERTIFICATE, password=CERT_PASSWORD, backend=default_backend()) pkb = p_key.private_bytes( encoding=serialization.Encoding.DER, format=serialization.PrivateFormat.PKCS8, encryption_algorithm=serialization.NoEncryption(), ) params["private_key"] = pkb else: # Username + Password set_provided(params, "authenticator", AUTHENTICATOR) set_provided(params, "password", PASSWORD) return params def row_to_incident(column_descriptions, row): # pylint: disable=W9014 """ Create incident from data returned by queried database in fetch_incidents parameter: (list) column_descriptions The metadata that describes the values for each column in the 'data' parameter parameter: (dict) row The row of data where each cell's key in the row is the name of the column to which it belongs returns: Incident Object """ incident = {} occurred = row.get(DATETIME_COLUMN) timestamp = None if occurred: if isinstance(occurred, dttime | timedelta): err_msg = "The datetime field specified in the integration parameters must " err_msg += 'contain values of type "datetime" or "date".' raise Exception(err_msg) timestamp = occurred.timestamp() * 1000 else: err_msg = "Nothing found when trying to fetch the datetime field specified in" err_msg += " the integration parameters. Please check that the name was correct." err_msg += " If the field name was correct, verify that the returned value for" err_msg += " the specified field is not NULL for ALL of the rows to be fetched." raise Exception(err_msg) # Incident Title if INCIDENT_NAME_COLUMN: name = row.get(INCIDENT_NAME_COLUMN) else: name = "Snowflake Incident -- " name += convert_datetime_to_string(occurred) + "- " + str(datetime.now().timestamp()) incident["name"] = name incident["occurred"] = timestamp_to_datestring(timestamp) # Incident occurrence time as timestamp - the datetime field specified in the integration parameters incident["timestamp"] = timestamp # The raw response for the row (reformatted to be json serializable) returned by the db query reformatted_row = format_to_json_serializable(column_descriptions, row) incident["rawJSON"] = json.dumps(reformatted_row) return incident """MAIN FUNCTIONS / API CALLS""" def test_module(): """ Test the validity of the integration instance parameters by trying to create a connection returns: An 'ok' message if valid, otherwise an error message """ params = get_connection_params({}) with snowflake.connector.connect(**params): # pylint: disable=E1101 demisto.results("ok") # pylint: disable=W9008 def fetch_incidents(): """ Fetch events from this integration and return them as Demisto incidents returns: Demisto incidents """ # demisto.getLastRun() will returns an obj with the previous run in it. last_run = demisto.getLastRun() # Get the last fetch time and data if it exists last_fetch = last_run.get("last_fetched_data_timestamp") last_fetched_data = last_run.get("last_fetched_data") # Handle first time fetch, fetch incidents retroactively if not last_fetch: last_fetch, _ = parse_date_range(FETCH_TIME, to_timestamp=True) args = {"limit": MAX_ROWS, "query": FETCH_QUERY} column_descriptions, data = snowflake_query(args) data.sort(key=lambda k: k[DATETIME_COLUMN]) # convert the data/events to demisto incidents incidents = [] for row in data: incident = row_to_incident(column_descriptions, row) incident_timestamp = incident.get("timestamp") # Update last run and add incident if the incident is newer than last fetch if incident_timestamp and incident_timestamp >= last_fetch: last_fetch = incident_timestamp if incident.get("rawJSON") != last_fetched_data: last_fetched_data = incident.get("rawJSON") del incident["timestamp"] incidents.append(incident) this_run = {"last_fetched_data": last_fetched_data, "last_fetched_data_timestamp": last_fetch} demisto.setLastRun(this_run) demisto.incidents(incidents) def snowflake_query(args): # pylint: disable=W9014 params = get_connection_params(args) query = args.get("query") limit = args.get("limit", "100") try: limit = int(limit) except ValueError: raise ValueError("The value for limit must be an integer.") if limit > MAX_ROWS: limit = MAX_ROWS with snowflake.connector.connect(**params) as connection, connection.cursor(snowflake.connector.DictCursor) as cur: # pylint: disable=E1101 cur.execute(query) results = cur.fetchmany(limit) if results: return cur.description, results else: return [], [] def snowflake_query_command(): args = demisto.args() # pylint: disable=W9017 query = args.get("query") db = args.get("database") if args.get("database") else DATABASE schema = args.get("schema") if args.get("schema") else SCHEMA col_descriptions, results = snowflake_query(args) if not results: demisto.results("No data found matching the query") # pylint: disable=W9008 else: results = format_to_json_serializable(col_descriptions, results) entry_context = {"Database": db, "Schema": schema, "Query": query, "Result": results} columns = argToList(args.get("columns")) human_readable = tableToMarkdown(query, results, columns, removeNull=True) demisto_transform = "Snowflake(val.Query && val.Query === obj.Query" demisto_transform += " && val.Database && val.Database === obj.Database" demisto_transform += " && val.Schema && val.Schema === obj.Schema)" outputs = {demisto_transform: entry_context} return_outputs(outputs=outputs, readable_output=human_readable, raw_response=results) # pylint: disable=W9009 def snowflake_update_command(): args = demisto.args() # pylint: disable=W9017 db_operation = args.get("db_operation") params = get_connection_params(args) with snowflake.connector.connect(**params) as connection, connection.cursor() as cursor: # pylint: disable=E1101 cursor.execute(db_operation) demisto.results("Operation executed successfully.") # pylint: disable=W9008 """COMMAND SWITCHBOARD""" commands = { "test-module": test_module, "fetch-incidents": fetch_incidents, "snowflake-query": snowflake_query_command, "snowflake-update": snowflake_update_command, } """EXECUTION""" def main() -> None: error = "" try: handle_proxy() if demisto.command() in commands: commands[demisto.command()]() except snowflake.connector.errors.Error as e: # pylint: disable=E1101 error = error_message_from_snowflake_error(e) except Exception as e: if IS_FETCH: raise e else: if isinstance(e, snowflake.connector.errors.Error): # pylint: disable=E1101 error = error_message_from_snowflake_error(e) else: error = str(e) finally: if error: return_error(error) if __name__ in ("__main__", "__builtin__", "builtins"): main()