DBotPredictOutOfTheBoxV2

Predict phishing incidents using the out-of-the-box pre-trained model.

python · Machine Learning

Source

# pylint: disable=no-member
import traceback

import demisto_ml
from CommonServerPython import *

TARGET_PRECISION = 0.97
THRESHOLD = 0.9
OUT_OF_THE_BOX_MODEL_NAME = "demisto_out_of_the_box_model_v2"
OUT_OF_THE_BOX_MODEL_PATH = "/ml/encrypted_model.b"
EVALUATION_PATH = "/ml/oob_evaluation.txt"
OOB_VERSION_INFO_KEY = "oob_version"


def oob_model_exists_and_updated():
    res_model = demisto.executeCommand("getMLModel", {"modelName": OUT_OF_THE_BOX_MODEL_NAME})[0]
    if is_error(res_model):
        return False
    model_type = dict_safe_get(res_model, [0, "Contents", "model", "type", "type"], "UNKNOWN_MODEL_TYPE")
    return model_type == demisto_ml.ModelType.Torch.value


def load_oob_model():
    try:
        encoded_model = demisto_ml.load_oob(OUT_OF_THE_BOX_MODEL_PATH)
    except Exception:
        return_error(traceback.format_exc())
    res = demisto.executeCommand(
        "createMLModel",
        {
            "modelData": encoded_model,
            "modelName": OUT_OF_THE_BOX_MODEL_NAME,
            "modelLabels": ["Malicious", "Non-Malicious"],
            "modelOverride": "true",
            "modelType": demisto_ml.ModelType.Torch.value,
            "modelExtraInfo": {"threshold": THRESHOLD},
        },
    )
    if is_error(res):
        return_error(get_error(res))

    with open(EVALUATION_PATH) as json_file:
        data = json.load(json_file)
    y_test = data["YTrue"]
    y_pred = data["YPred"]
    y_pred_prob = data["YPredProb"]

    y_pred_evaluation = [{pred: prob} for pred, prob in zip(y_pred, y_pred_prob)]
    res = demisto.executeCommand(
        "GetMLModelEvaluation",
        {
            "yTrue": json.dumps(y_test),
            "yPred": json.dumps(y_pred_evaluation),
            "targetPrecision": str(0.85),
            "targetRecall": str(0),
            "detailedOutput": "true",
        },
    )
    if is_error(res):
        return_error(get_error(res))
    confusion_matrix = json.loads(res[0]["Contents"]["csr_matrix_at_threshold"])
    confusion_matrix_no_all = {k: v for k, v in confusion_matrix.items() if k != "All"}
    confusion_matrix_no_all = {
        k: {sub_k: sub_v for sub_k, sub_v in v.items() if sub_k != "All"} for k, v in confusion_matrix_no_all.items()
    }
    res = demisto.executeCommand(
        "evaluateMLModel",
        {
            "modelConfusionMatrix": confusion_matrix_no_all,
            "modelName": OUT_OF_THE_BOX_MODEL_NAME,
            "modelEvaluationVectors": {"Ypred": y_pred, "Ytrue": y_test, "YpredProb": y_pred_prob},
            "modelConfidenceThreshold": THRESHOLD,
            "modelTargetPrecision": TARGET_PRECISION,
        },
    )

    if is_error(res):
        return_error(get_error(res))


def predict_phishing_words():
    if not oob_model_exists_and_updated():
        load_oob_model()
    dargs = demisto.args()
    dargs["modelName"] = OUT_OF_THE_BOX_MODEL_NAME
    dargs["modelStoreType"] = "mlModel"
    res = demisto.executeCommand("DBotPredictPhishingWords", dargs)
    if is_error(res):
        return_error(get_error(res))
    return res


def main():
    res = predict_phishing_words()
    return res


if __name__ in ["__main__", "__builtin__", "builtins"]:
    demisto.results(main())

README

Predict phishing incidents using the out-of-the-box pre-trained model.

Script Data

Name Description
Script Type python3
Tags phishing, ml
Cortex XSOAR Version 5.5.0

Inputs

Argument Name Description
emailSubject Subject of the email.
emailBody Body of the email.
emailBodyHTML HTML body of the email. Only use this field if the emailBody argument is empty.
topWordsLimit Maximum number of positive/negative words to return for the model decision.
wordThreshold Threshold to determine word importance (range 0-1). Default is 0.05.
minTextLength Minimum number of characters for the prediction.
labelProbabilityThreshold The label probability threshold. Default is 0.
confidenceThreshold The confidence threshold. The model will provide predictions only if their confidence is above this threshold.
returnError Whether to return an error when there is no prediction. Default is “true”.
setIncidentFields Whether to set Cortex XSOAR out-of-the-box DBot fields.

Outputs

Path Description Type
DBotPredictPhishingWords.Label The predicted label. String
DBotPredictPhishingWords.Probability The predicted probability (range 0-1). Number
DBotPredictPhishingWords.PositiveWords A list of words in the input text that supports the model decision. Unknown
DBotPredictPhishingWords.NegativeWords A list of words in the input text that do not support the model decision. These words better support a different classification class. Unknown
DBotPredictPhishingWords.TextTokensHighlighted The input text (after pre-processing) with the positive words that support the model decision. String