import numpy as np
import pandas as pd
from CommonServerPython import *
ALL_LABELS = "*"
PREDICTIONS_OUT_FILE_NAME = "predictions.csv"
def canonize_label(label):
return label.replace(" ", "_")
def get_phishing_map_labels(comma_values):
if comma_values == ALL_LABELS:
return comma_values
values = [x.strip() for x in comma_values.split(",")]
labels_dict = {}
for v in values:
v = v.strip()
if ":" in v:
splited = v.rsplit(":", maxsplit=1)
labels_dict[splited[0].strip()] = splited[1].strip()
else:
labels_dict[v] = v
if len(set(labels_dict.values())) == 1:
mapped_value = list(labels_dict.values())[0]
error = [f"Label mapping error: you need to map to at least two labels: {mapped_value}."]
return_error("\n".join(error))
return {k: canonize_label(v) for k, v in labels_dict.items()}
def get_data_with_mapped_label(y_true_list, labels_mapping):
mapped_y_true = []
relevant_indices = []
for i, y_true in enumerate(y_true_list):
if labels_mapping == ALL_LABELS:
mapped_y_true.append(canonize_label(y_true))
relevant_indices.append(i)
elif y_true in labels_mapping:
mapped_y_true.append(canonize_label(labels_mapping[y_true]))
relevant_indices.append(i)
else:
continue
return mapped_y_true, relevant_indices
def get_ml_model_evaluation(y_test, y_pred, target_accuracy, target_recall, detailed=False):
res = demisto.executeCommand(
"GetMLModelEvaluation",
{
"yTrue": json.dumps(y_test),
"yPred": json.dumps(y_pred),
"targetPrecision": str(target_accuracy),
"targetRecall": str(target_recall),
"detailedOutput": "true" if detailed else "false",
},
)
if is_error(res):
return_error(get_error(res))
return res
def output_model_evaluation(y_test, y_pred, res, context_field, human_readable_title=None):
threshold = float(res[0]["Contents"]["threshold"])
confusion_matrix = json.loads(res[0]["Contents"]["csr_matrix_at_threshold"])
metrics_df = json.loads(res[0]["Contents"]["metrics_df"])
human_readable = res[0]["HumanReadable"]
if human_readable_title is not None:
human_readable = "\n".join([human_readable_title, human_readable])
result_entry = {
"Type": entryTypes["note"],
"Contents": {
"Threshold": threshold,
"ConfusionMatrixAtThreshold": confusion_matrix,
"Metrics": metrics_df,
"YTrue": y_test,
"YPred": y_pred,
},
"ContentsFormat": formats["json"],
"HumanReadable": human_readable,
"HumanReadableFormat": formats["markdown"],
"EntryContext": {
context_field: {
"EvaluationScores": metrics_df,
"ConfusionMatrix": confusion_matrix,
}
},
}
demisto.results(result_entry)
return confusion_matrix
def return_file_result_with_predictions_on_test_set(data, y_true, y_pred, y_pred_prob, additional_fields):
predictions_data = {}
for field in additional_fields:
predictions_data[field] = [i.get(field, "") for i in data]
predictions_data["y_true"] = y_true
predictions_data["y_pred"] = y_pred
predictions_data["y_pred_prob"] = y_pred_prob
df = pd.DataFrame(predictions_data)
non_empty_columns = [field for field in additional_fields if df[field].astype(bool).any()]
csv_df = df.to_csv(columns=["y_true", "y_pred", "y_pred_prob"] + non_empty_columns, encoding="utf-8")
demisto.results(fileResult(PREDICTIONS_OUT_FILE_NAME, csv_df))
def main(
incident_types,
incident_query,
y_true_field,
y_pred_field,
y_pred_prob_field,
model_target_accuracy,
labels_mapping,
additional_fields,
):
non_empty_fields = f"{y_true_field.strip()},{y_pred_field.strip()}"
incidents_query_args = {
"incidentTypes": incident_types,
"NonEmptyFields": non_empty_fields,
}
if incident_query is not None:
incidents_query_args["query"] = incident_query
incidents_query_res = demisto.executeCommand("GetIncidentsByQuery", incidents_query_args)
if is_error(incidents_query_res):
return_error(get_error(incidents_query_res))
incidents = json.loads(incidents_query_res[0]["Contents"])
if incidents:
demisto.results(f"Found {len(incidents)} incident(s)")
y_true = []
y_pred = []
y_pred_prob = []
incidents_with_missing_pred_prob = 0
for i in incidents:
y_true.append(i[y_true_field])
y_pred.append(i[y_pred_field])
if y_pred_prob_field not in i:
incidents_with_missing_pred_prob += 1
y_pred_prob.append(i.get(y_pred_prob_field, None))
y_true, relevant_indices = get_data_with_mapped_label(y_true, labels_mapping)
y_pred = [y_pred[i] for i in relevant_indices]
y_pred_prob = [y_pred_prob[i] for i in relevant_indices]
incidents = [incidents[i] for i in relevant_indices]
y_pred_prob_is_given = incidents_with_missing_pred_prob == 0
if y_pred_prob_is_given:
y_pred_dict = [{label: prob} for label, prob in zip(y_pred, y_pred_prob)]
else:
y_pred_dict = [{label: 1.0} for label in y_pred]
if y_pred_prob_is_given:
res_threshold = get_ml_model_evaluation(y_true, y_pred_dict, model_target_accuracy, target_recall=0, detailed=True)
# show results for the threshold found - last result so it will appear first
output_model_evaluation(
y_test=y_true, y_pred=y_pred_dict, res=res_threshold, context_field="EvaluateMLModllAtProduction"
)
# show results if no threshold (threhsold=0) was used. Following code is reached only if a legal thresh was found:
if not y_pred_prob_is_given or not np.isclose(float(res_threshold[0]["Contents"]["threshold"]), 0):
res = get_ml_model_evaluation(y_true, y_pred_dict, target_accuracy=0, target_recall=0)
human_readable = "\n".join(
["## Results for No Threshold", "The following results were achieved by using no threshold (threshold equals 0)"]
)
output_model_evaluation(
y_test=y_true,
y_pred=y_pred_dict,
res=res,
context_field="EvaluateMLModllAtProductionNoThresh",
human_readable_title=human_readable,
)
return_file_result_with_predictions_on_test_set(incidents, y_true, y_pred, y_pred_prob, additional_fields)
else:
return_results("No incidents found.")
model_target_accuracy = demisto.args().get("modelTargetAccuracy", 0)
incident_types = demisto.args()["incidentTypes"]
incident_query = demisto.args().get("incidentsQuery", None)
y_true_field = demisto.args()["emailTagKey"]
y_pred_field = demisto.args()["emailPredictionKey"]
y_pred_prob_field = demisto.args()["emailPredictionProbabilityKey"]
labels_mapping = get_phishing_map_labels(demisto.args()["phishingLabels"])
additional_fields = demisto.args().get("additionalFields", "")
additional_fields = additional_fields.split(",")
additional_fields = [x.strip() for x in additional_fields]
main(
incident_types,
incident_query,
y_true_field,
y_pred_field,
y_pred_prob_field,
model_target_accuracy,
labels_mapping,
additional_fields,
)