# pylint: disable=no-member import gc from collections import Counter, defaultdict import demisto_ml import pandas as pd from CommonServerPython import * from sklearn.model_selection import StratifiedKFold ALL_LABELS = "*" GENERAL_SCORES = { "micro avg": "The metrics is applied globally by counting the total true positives, false negatives and false positives", "macro avg": "The metrics is applied for each label, and find their unweighted mean.", "weighted avg": "The metrics is applied for each label, and find their average weighted by support " "(the number of true instances for each label). This alters macro to account for label imbalance;", } DBOT_TAG_FIELD = "dbot_internal_tag_field" MIN_INCIDENTS_THRESHOLD = 100 PREDICTIONS_OUT_FILE_NAME = "predictions_on_test_set.csv" # FROM_SCRATCH_TRAINING_ALGO is the UI equivalent of FASTTEXT_TRAINING_ALGO FROM_SCRATCH_TRAINING_ALGO = "from_scratch" FINETUNE_TRAINING_ALGO = "fine_tune" FASTTEXT_TRAINING_ALGO = "fasttext" AUTO_TRAINING_ALGO = "auto" # the following mapping need to correspond to predict_phishing_words func at DBotPredictPhishingWords ALGO_TO_MODEL_TYPE = {FASTTEXT_TRAINING_ALGO: "fasttext_phishing", FINETUNE_TRAINING_ALGO: "torch_phishing"} FINETUNE_LABELS = ["Malicious", "Non-Malicious"] 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] return_error(f"Label mapping error: you need to map to at least two labels: {mapped_value}.") return {k.encode("utf-8", "ignore").decode("utf-8"): v for k, v in labels_dict.items()} def read_file(input_data, input_type): data = [] # type: List[Dict[str, str]] file_path, file_content = "", "" if not input_data: return data if input_type.endswith("string"): if "b64" in input_type: input_data = base64.b64decode(input_data) file_content = input_data.decode("utf-8") else: file_content = input_data else: res = demisto.getFilePath(input_data) if not res: return_error(f"Entry {input_data} not found") file_path = res["path"] if input_type.startswith("json"): with open(file_path) as f: file_content = f.read() if input_type.startswith("csv"): return pd.read_csv(file_path).fillna("").to_dict(orient="records") elif input_type.startswith("json"): return json.loads(file_content) elif input_type.startswith("pickle"): return pd.read_pickle(file_path, compression=None) else: return_error(f"Unsupported file type {input_type}") return None def get_file_entry_id(file_name): file_name = file_name.strip() res = demisto.dt(demisto.context(), f"File(val.Name == '{file_name}')") if not res or len(res) == 0: return_error(f"Cannot find file entry id in context by filename: {file_name}") if isinstance(res, list): res = res[0] return res["EntryID"] def read_files_by_name(file_names, input_type): names = filter(None, argToList(file_names)) # type: ignore[var-annotated] data = [] for name in names: data += read_file(get_file_entry_id(name), input_type) return data def get_data_with_mapped_label(data, labels_mapping, tag_field): new_data = [] exist_labels_counter = defaultdict(int) # type: Dict[str, int] missing_labels_counter = defaultdict(int) # type: Dict[str, int] for row in data: original_label = row[tag_field] if labels_mapping == ALL_LABELS: row[tag_field] = original_label else: if original_label in labels_mapping: row[tag_field] = labels_mapping[original_label] elif original_label.lower() in labels_mapping: original_label = original_label.lower() row[tag_field] = labels_mapping[original_label] else: missing_labels_counter[original_label] += 1 continue exist_labels_counter[original_label] += 1 new_data.append(row) return new_data, dict(exist_labels_counter), dict(missing_labels_counter) def store_model_in_demisto( model_name, model_override, X, y, confusion_matrix, threshold, y_test_true, y_test_pred, y_test_pred_prob, target_accuracy, algorithm, ): global ALGO_TO_MODEL_TYPE phishing_model = demisto_ml.train_model_handler(X, y, algorithm=algorithm, compress=True) model_labels = phishing_model.get_model_labels() model_data = phishing_model.dumps() res = demisto.executeCommand( "createMLModel", { "modelData": model_data, "modelName": model_name, "modelLabels": model_labels, "modelOverride": model_override, "modelExtraInfo": {"threshold": threshold}, "modelType": ALGO_TO_MODEL_TYPE[algorithm], }, ) if is_error(res): return_error(get_error(res)) y_test_pred_prob = [float(x) for x in y_test_pred_prob] res = demisto.executeCommand( "evaluateMLModel", { "modelConfusionMatrix": confusion_matrix, "modelName": model_name, "modelEvaluationVectors": {"Ypred": y_test_pred, "Ytrue": y_test_true, "YpredProb": y_test_pred_prob}, "modelConfidenceThreshold": threshold, "modelTargetPrecision": target_accuracy, }, ) if is_error(res): return_error(get_error(res)) def find_keywords(data, tag_field, text_field, min_score): keywords = demisto_ml.get_keywords_for_labels(data, tag_field, text_field) human_readable = "# Keywords per category\n" for category, scores in keywords.items(): sorted_scores = sorted(scores.items(), key=lambda x: x[1], reverse=True) table_items = [{"Word": word, "Score": f"{score:.2f}"} for word, score in sorted_scores if score >= min_score] human_readable += tableToMarkdown(category, table_items, ["Word", "Score"]) demisto.results( { "Type": entryTypes["note"], "Contents": keywords, "ContentsFormat": formats["json"], "HumanReadable": human_readable, "HumanReadableFormat": formats["markdown"], } ) def set_tag_field(data, tag_fields): empty_label_indices = [] for i, d in enumerate(data): found_field = False for field in tag_fields: if d.get(field) is not None: label = d[field] if isinstance(label, list) and len(label) > 0: label = label[0] elif isinstance(label, list) and len(label) == 0: continue label = label.encode("utf-8", "ignore").decode("utf-8") d[DBOT_TAG_FIELD] = str(label) found_field = True break if not found_field: empty_label_indices.append(i) data = [d for i, d in enumerate(data) if i not in empty_label_indices] return data def output_model_evaluation(model_name, y_test, y_pred, res, context_field, human_readable_title=None): threshold = float(res["Contents"]["threshold"]) confusion_matrix_at_thresh = json.loads(res["Contents"]["csr_matrix_at_threshold"]) confusion_matrix_no_thresh = json.loads(res["Contents"]["csr_matrix_no_threshold"]) metrics_df = json.loads(res["Contents"]["metrics_df"]) human_readable = res["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_at_thresh, "ConfusionMatrixNoThreshold": confusion_matrix_no_thresh, "Metrics": metrics_df, "YTrue": y_test, "YPred": y_pred, }, "ContentsFormat": formats["json"], "HumanReadable": human_readable, "HumanReadableFormat": formats["markdown"], "EntryContext": { context_field: { "ModelName": model_name, "EvaluationScores": metrics_df, "ConfusionMatrix": confusion_matrix_at_thresh, "ConfusionMatrixNoThresh": confusion_matrix_no_thresh, } }, } demisto.results(result_entry) confusion_matrix_at_thresh = {k: v for k, v in confusion_matrix_at_thresh.items() if k != "All"} confusion_matrix_at_thresh = { k: {sub_k: sub_v for sub_k, sub_v in v.items() if sub_k != "All"} for k, v in confusion_matrix_at_thresh.items() } return confusion_matrix_at_thresh, metrics_df 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[0] def validate_data_and_labels(data, exist_labels_counter, labels_mapping, missing_labels_counter): labels_counter = Counter([x[DBOT_TAG_FIELD] for x in data]) labels_below_thresh = [label for label, count in labels_counter.items() if count < MIN_INCIDENTS_THRESHOLD] if len(labels_below_thresh) > 0: err = [f"Minimum number of incidents per label required for training is {MIN_INCIDENTS_THRESHOLD}."] err += [f"The following labels have less than {MIN_INCIDENTS_THRESHOLD} incidents: "] for x in labels_below_thresh: err += [f"- {x}: {labels_counter[x]!s}"] err += ["Make sure that enough incidents exist in the environment per each of these labels."] missing_labels = ", ".join(missing_labels_counter.keys()) err += [f"The following labels were not mapped to any label in the labels mapping: {missing_labels}."] if labels_mapping != ALL_LABELS: err += ["The given mapped labels are: {}.".format(", ".join(labels_mapping.keys()))] return_error("\n".join(err)) if len(exist_labels_counter) == 0: err = ["Did not found any incidents with labels of the labels mapping."] if len(missing_labels_counter) > 0: err += ["The following labels were found: {}".format(", ".join(k for k in missing_labels_counter))] err += ["Please include these labels at the mapping, or change the query to include your relevant labels"] return_error("\n".join(err)) if len(missing_labels_counter) > 0: human_readable = tableToMarkdown("Skip labels - did not match any of specified labels", missing_labels_counter) entry = { "Type": entryTypes["note"], "Contents": missing_labels_counter, "ContentsFormat": formats["json"], "HumanReadable": human_readable, "HumanReadableFormat": formats["markdown"], } demisto.results(entry) if len(exist_labels_counter) > 0: exist_labels_counter_mapped = {} for label, count in exist_labels_counter.items(): mapped_label = labels_mapping[label] if isinstance(labels_mapping, dict) else label if mapped_label != label: label = f"{label} -> {mapped_label}" exist_labels_counter_mapped[label] = count human_readable = tableToMarkdown("Found labels", exist_labels_counter_mapped) entry = { "Type": entryTypes["note"], "Contents": exist_labels_counter, "ContentsFormat": formats["json"], "HumanReadable": human_readable, "HumanReadableFormat": formats["markdown"], } demisto.results(entry) if len({x[DBOT_TAG_FIELD] for x in data}) == 1: single_label = [x[DBOT_TAG_FIELD] for x in data][0] if labels_mapping == ALL_LABELS: err = [f"All received incidents have the same label: {single_label}."] else: err = [f"All received incidents mapped to the same label: {single_label}."] err += ["At least 2 different labels are required to train a classifier."] if labels_mapping == ALL_LABELS: err += ["Please make sure that incidents of at least 2 labels exist in the environment."] else: err += ["The following labels were not mapped to any label in the labels mapping:"] err += [", ".join(list(missing_labels_counter))] not_found_mapped_label = [x for x in labels_mapping if x not in exist_labels_counter or exist_labels_counter[x] == 0] if len(not_found_mapped_label) > 0: miss = ", ".join(not_found_mapped_label) err += [f"Notice that the following mapped labels were not found among all incidents: {miss}."] return_error("\n".join(err)) def return_file_result_with_predictions_on_test_set(data, original_text_fields, test_index, text_field, y_test, y_pred_dict): if original_text_fields is None or original_text_fields.strip() == "": original_text_fields = [text_field] else: original_text_fields = re.split(r"[|,]", original_text_fields) original_text_fields = [x.strip() for x in original_text_fields] + [text_field] predictions_data = {} test_data = [data[i] for i in test_index] for field in original_text_fields: predictions_data[field] = [record.get(field, "") for record in test_data] predictions_data["y_true"] = y_test y_pred = [] y_pred_prob = [] for y_i in y_pred_dict: y_pred_prob_i = max(y_i.values()) y_pred_i = [label for label, label_prob in y_i.items() if y_i[label] == y_pred_prob_i][0] y_pred.append(y_pred_i) y_pred_prob.append(y_pred_prob_i) 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 original_text_fields if df[field].astype(bool).any()] csv_df = df.to_csv(columns=non_empty_columns + ["y_true", "y_pred", "y_pred_prob"], encoding="utf-8") demisto.results(fileResult(PREDICTIONS_OUT_FILE_NAME, csv_df)) def get_train_and_test_sets_indices(X, y): train_set_ratio = float(demisto.args()["trainSetRatio"]) n_splits = int(1.0 / (1 - train_set_ratio)) skf = StratifiedKFold(n_splits=n_splits, shuffle=False, random_state=None) skf.get_n_splits(X, y) train_index, test_index = list(skf.split(X, y))[-1] return test_index, train_index def get_X_and_y_from_data(data, text_field): y = [x[DBOT_TAG_FIELD] for x in data] X = [x[text_field] for x in data] if len(X) != len(y): return_error("Error: data and tag data are different length") return X, y def validate_labels_and_decide_algorithm(y, algorithm): labels_counter = Counter(y) # type: Dict[str, int] illegal_labels_for_fine_tune = [label for label in labels_counter if label not in FINETUNE_LABELS] if algorithm == FINETUNE_TRAINING_ALGO and len(illegal_labels_for_fine_tune) > 0: error = [ "When trainingAlgorithm is set to {}, all labels must be mapped to {}.\n".format( algorithm, ", ".join(FINETUNE_LABELS) ) ] error += ["The following labels/verdicts need to be mapped to one of those values: "] error += [", ".join(illegal_labels_for_fine_tune) + "."] return_error("\n".join(error)) return None elif algorithm == AUTO_TRAINING_ALGO: return FASTTEXT_TRAINING_ALGO else: return algorithm def validate_confusion_matrix(confusion_matrix): for label in confusion_matrix: tp = confusion_matrix[label][label] fp = sum(confusion_matrix[label_other][label] for label_other in confusion_matrix if label != label_other) if tp == fp == 0: return False return True def main(): input = demisto.args().get("input") input_type = demisto.args().get("inputType", "pickle_filename") model_name = demisto.args().get("modelName", "phishing_model") store_model = demisto.args().get("storeModel") == "true" model_override = demisto.args().get("overrideExistingModel", "false") == "true" target_accuracy = float(demisto.args().get("targetAccuracy", "0.8")) text_field = demisto.args().get("textField", "dbot_processed_text") tag_fields = demisto.args().get("tagField").split(",") labels_mapping = get_phishing_map_labels(demisto.args().get("phishingLabels")) keyword_min_score = float(demisto.args().get("keywordMinScore", "0.05")) return_predictions_on_test_set = demisto.args().get("returnPredictionsOnTestSet", "false") == "true" original_text_fields = demisto.args().get("originalTextFields", "emailsubject|name,emailbody|emailbodyhtml") algorithm = demisto.args().get("trainingAlgorithm", AUTO_TRAINING_ALGO) # FASTTEXT_TRAINING_ALGO and FROM_SCRATCH_TRAINING_ALGO are equivalent, replacement is done because ml_lib # expects algorithm as one of (FASTTEXT_TRAINING_ALGO, FINETUNE_TRAINING_ALGO) algorithm = FASTTEXT_TRAINING_ALGO if algorithm == FROM_SCRATCH_TRAINING_ALGO else algorithm if input_type.endswith("filename"): data = read_files_by_name(input, input_type.split("_")[0].strip()) else: data = read_file(input, input_type) if len(data) < MIN_INCIDENTS_THRESHOLD: return_results( f"{len(data)} incident(s) received." f"\nMinimum number of incidents per label required for training: {MIN_INCIDENTS_THRESHOLD}." "\nMake sure that all arguments are set correctly and that enough incidents exist in the environment." ) else: data = set_tag_field(data, tag_fields) data, exist_labels_counter, missing_labels_counter = get_data_with_mapped_label(data, labels_mapping, DBOT_TAG_FIELD) validate_data_and_labels(data, exist_labels_counter, labels_mapping, missing_labels_counter) # print important words for each category find_keywords_bool = "findKeywords" in demisto.args() and demisto.args()["findKeywords"] == "true" if find_keywords_bool: try: find_keywords(data, DBOT_TAG_FIELD, text_field, keyword_min_score) except Exception: pass X, y = get_X_and_y_from_data(data, text_field) algorithm = validate_labels_and_decide_algorithm(y, algorithm) test_index, train_index = get_train_and_test_sets_indices(X, y) X_train, X_test = [X[i] for i in train_index], [X[i] for i in test_index] y_train, y_test = [y[i] for i in train_index], [y[i] for i in test_index] phishing_model = demisto_ml.train_model_handler(X_train, y_train, algorithm=algorithm, compress=False) ft_test_predictions = phishing_model.predict(X_test) y_pred = [{y_tuple[0]: float(y_tuple[1])} for y_tuple in ft_test_predictions] if return_predictions_on_test_set: return_file_result_with_predictions_on_test_set(data, original_text_fields, test_index, text_field, y_test, y_pred) target_recall = 1 - float(demisto.args().get("maxBelowThreshold", 1)) threshold_metrics_entry = get_ml_model_evaluation(y_test, y_pred, target_accuracy, target_recall, detailed=True) # show results for the threshold found - last result so it will appear first confusion_matrix, metrics_json = output_model_evaluation( model_name=model_name, y_test=y_test, y_pred=y_pred, res=threshold_metrics_entry, context_field="DBotPhishingClassifier", ) actual_min_accuracy = min(v for k, v in metrics_json["Precision"].items() if k != "All") if store_model: del phishing_model gc.collect() if not validate_confusion_matrix(confusion_matrix): return_error( "The trained model didn't manage to predict some of the classes. This model won't be stored." "Please try to retrain the model using a different configuration." ) y_test_pred = [y_tuple[0] for y_tuple in ft_test_predictions] y_test_pred_prob = [y_tuple[1] for y_tuple in ft_test_predictions] threshold = float(threshold_metrics_entry["Contents"]["threshold"]) store_model_in_demisto( model_name=model_name, model_override=model_override, X=X, y=y, confusion_matrix=confusion_matrix, threshold=threshold, y_test_true=y_test, y_test_pred=y_test_pred, y_test_pred_prob=y_test_pred_prob, target_accuracy=actual_min_accuracy, algorithm=algorithm, ) demisto.results(f"Done training on {len(y)} samples model stored successfully") else: demisto.results("Skip storing model") if __name__ in ["builtins", "__main__"]: main()