DBotPredictPhishingWords

Predict text label using a pre-trained machine learning phishing model, and get the most important words used in the classification decision.

python · Base

Details

ID	`DBotPredictPhishingWords`
Language	python
From Version	5.0.0
Docker Image	`demisto/ml:1.0.0.9793547`
Tags	ml phishing

README

Predict text label using a pre-trained machine learning phishing model, and get the most important words used in the classification decision.

Script Data

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

Used In

This script is used in the following playbooks and scripts.

Phishing Investigation - Generic v2

Inputs

Argument Name	Description
modelName	The model name (or list name) in Demisto.
hashSeed	Seed for the hash function, at the pre-process stage.
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. Default is 20.
wordThreshold	Threshold to determine word importance (range 0-1). Default is 0.05.
modelStoreType	How the model is stored in Demisto. Can be “list” or “mlModel”. Default is “list”.
minTextLength	Minimum number of characters for the prediction.
labelProbabilityThreshold	The label probability threshold. Default is 0.8.
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 Demisto out-of-the-box DBot fields.
language	The language of the input text. Default is “Any”. Can be “Any”, “English”, “German”, “French”, “Spanish”, “Portuguese”, “Italian”, “Dutch”, or “Other”. If “Any” or “Other” is selected, the script preprocess the entire input, no matter what its acutual language is. If a specific language is selected, the script filters out any other language from the output text.
tokenizationMethod	Tokenization method for text. Only required when the language argument is set to “Other”. Can be “tokenizer”, “byWords”, or “byLetters”.

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

Source

DBotPredictPhishingWords.py DBotPredictPhishingWords_test.py DBotPredictPhishingWords.yml

Download History Raw

from collections import defaultdict

import pytest
from CommonServerPython import *
from DBotPredictPhishingWords import get_model_data, main, predict_phishing_words

TOKENIZATION_RESULT = None


class PhishingModelMock:
    def __init__(self, filter_words_res=None, explain_model_words_res=None):
        self.filter_words_res = filter_words_res
        self.explain_model_words_res = explain_model_words_res

    def filter_model_words(self):
        return self.filter_words_res

    def explain_model_words(self, a, b, c, d):
        return self.explain_model_words_res


def get_args():
    args = defaultdict(lambda: "yes")
    args["encoding"] = "utf8"
    args["encoding"] = "utf8"
    args["removeNonEnglishWords"] = "no"
    return args


def bold(word):
    return f"**{word}**"


def executeCommand(command, args=None):
    global TOKENIZATION_RESULT
    if command == "getList":
        return [{"Contents": "ModelDataList", "Type": "note"}]
    elif command == "getMLModel":
        return [{"Contents": {"modelData": "ModelDataML", "model": {"type": {"type": ""}}}, "Type": "note"}]
    elif command == "DBotPreProcessTextData":
        TOKENIZATION_RESULT["originalText"] = args["input"]
        TOKENIZATION_RESULT["tokenizedText"] = args["input"]
        return [{"Contents": TOKENIZATION_RESULT, "Type": "note"}]
    elif command == "HighlightWords":
        text = args["text"]
        terms = set(args["terms"].split(","))
        words = text.split()
        for i, w in enumerate(words):
            if w in terms:
                words[i] = bold(w)
        return [{"Contents": " ".join(words), "Type": "note"}]
    return None


def test_get_model_data(mocker):
    mocker.patch.object(demisto, "executeCommand", side_effect=executeCommand)
    assert get_model_data("test", "list", True)[0] == "ModelDataList"
    assert get_model_data("test", "mlModel", True)[0] == "ModelDataML"


def test_predict_phishing_words(mocker):
    global TOKENIZATION_RESULT
    d = {"Label": "Valid", "Probability": 0.7, "PositiveWords": ["word1"], "NegativeWords": ["word2"]}
    phishing_mock = PhishingModelMock()
    mocker.patch.object(demisto, "executeCommand", side_effect=executeCommand)
    mocker.patch.object(demisto, "args", return_value={"topWordsLimit": 10})
    mocker.patch("demisto_ml.phishing_model_loads_handler", return_value=phishing_mock, create=True)
    mocker.patch.object(demisto, "incidents", return_value=[{"isPlayground": True}])
    mocker.patch.object(phishing_mock, "filter_model_words", return_value=("text", 2), create=True)
    mocker.patch.object(phishing_mock, "explain_model_words", return_value=d, create=True)

    TOKENIZATION_RESULT = {
        "originalText": "word1 word2 word3",
        "tokenizedText": "word1 word2 word3",
        "originalWordsToTokens": {"word1": ["word1"], "word2": ["word2"], "word3": ["word3"]},
    }

    email_subject = "word1"
    email_body = "word2 word3"
    res = predict_phishing_words("modelName", "list", email_subject, email_body, 0, 0, 0, 10, True)
    correct_res = {
        "OriginalText": concatenate_subject_body(email_subject, email_body),
        "Probability": 0.7,
        "NegativeWords": ["word2"],
        "TextTokensHighlighted": concatenate_subject_body(f"**{email_subject}**", email_body),
        "PositiveWords": ["word1"],
        "Label": "Valid",
    }
    assert res["Contents"] == correct_res


def concatenate_subject_body(email_subject, email_body):
    return f"{email_subject} \n{email_body}"


def test_predict_phishing_words_low_threshold(mocker):
    global TOKENIZATION_RESULT
    d = {"Label": "Valid", "Probability": 0.7, "PositiveWords": ["word1"], "NegativeWords": ["word2"]}
    phishing_mock = PhishingModelMock()
    mocker.patch.object(demisto, "executeCommand", side_effect=executeCommand)
    mocker.patch.object(demisto, "args", return_value={"topWordsLimit": 10})
    mocker.patch("demisto_ml.phishing_model_loads_handler", return_value=phishing_mock, create=True)
    mocker.patch.object(phishing_mock, "filter_model_words", return_value=("text", 2), create=True)
    mocker.patch.object(phishing_mock, "explain_model_words", return_value=d, create=True)
    TOKENIZATION_RESULT = {
        "originalText": "word1 word2 word3",
        "tokenizedText": "word1 word2 word3",
        "hashedTokenizedText": "23423 432432 12321",
        "originalWordsToTokens": {"word1": ["word1"], "word2": ["word2"], "word3": ["word3"]},
        "wordsToHashedTokens": {"word1": ["23423"], "word2": ["432432"], "word3": ["12321"]},
    }
    with pytest.raises(SystemExit):
        predict_phishing_words("modelName", "list", "subject", "body", 0, 0.8, 0, 10, True)


def test_predict_phishing_words_no_words(mocker):
    global TOKENIZATION_RESULT
    phishing_mock = PhishingModelMock()
    mocker.patch.object(demisto, "executeCommand", side_effect=executeCommand)
    mocker.patch.object(demisto, "args", return_value={"topWordsLimit": 10})
    mocker.patch("demisto_ml.phishing_model_loads_handler", return_value=phishing_mock, create=True)
    mocker.patch.object(phishing_mock, "filter_model_words", return_value=("", 0), create=True)
    TOKENIZATION_RESULT = {
        "originalText": "word1 word2 word3",
        "tokenizedText": "word1 word2 word3",
        "hashedTokenizedText": "23423 432432 12321",
        "originalWordsToTokens": {"word1": ["word1"], "word2": ["word2"], "word3": ["word3"]},
        "wordsToHashedTokens": {"word1": ["23423"], "word2": ["432432"], "word3": ["12321"]},
    }
    with pytest.raises(SystemExit):
        predict_phishing_words("modelName", "list", "subject", "body", 0, 0, 0, 10, True)
    mocker.patch.object(phishing_mock, "filter_model_words", return_value=("", 10), create=True)
    with pytest.raises(SystemExit):
        predict_phishing_words("modelName", "list", "subject", "body", 20, 0, 0, 10, True)


def test_predict_phishing_words_hashed(mocker):
    global TOKENIZATION_RESULT
    d = {"Label": "Valid", "Probability": 0.7, "PositiveWords": ["23423"], "NegativeWords": ["432432"]}
    phishing_mock = PhishingModelMock()
    mocker.patch.object(demisto, "incidents", return_value=[{"isPlayground": True}])

    mocker.patch.object(demisto, "executeCommand", side_effect=executeCommand)
    mocker.patch.object(demisto, "args", return_value={"topWordsLimit": 10, "hashSeed": 10})
    mocker.patch("demisto_ml.phishing_model_loads_handler", return_value=phishing_mock, create=True)
    mocker.patch.object(phishing_mock, "filter_model_words", return_value=("text", 2), create=True)
    mocker.patch.object(phishing_mock, "explain_model_words", return_value=d, create=True)
    TOKENIZATION_RESULT = {
        "originalText": "word1 word2 word3",
        "tokenizedText": "word1 word2 word3",
        "hashedTokenizedText": "23423 432432 12321",
        "originalWordsToTokens": {"word1": ["word1"], "word2": ["word2"], "word3": ["word3"]},
        "wordsToHashedTokens": {"word1": ["23423"], "word2": ["432432"], "word3": ["12321"]},
    }
    email_subject = "word1"
    email_body = "word2 word3"
    res = predict_phishing_words("modelName", "list", email_subject, email_body, 0, 0, 0, 10, True)
    assert res["Contents"] == {
        "OriginalText": concatenate_subject_body(email_subject, email_body),
        "Probability": 0.7,
        "NegativeWords": ["word2"],
        "TextTokensHighlighted": concatenate_subject_body(f"**{email_subject}**", email_body),
        "PositiveWords": ["word1"],
        "Label": "Valid",
    }


def test_predict_phishing_words_tokenization_by_character(mocker):
    global TOKENIZATION_RESULT
    phishing_mock = PhishingModelMock()
    mocker.patch.object(demisto, "executeCommand", side_effect=executeCommand)
    mocker.patch.object(demisto, "args", return_value={"topWordsLimit": 10, "hashSeed": 10})
    mocker.patch("demisto_ml.phishing_model_loads_handler", return_value=phishing_mock, create=True)
    mocker.patch.object(phishing_mock, "filter_model_words", return_value=("text", 2), create=True)
    mocker.patch.object(demisto, "incidents", return_value=[{"isPlayground": True}])
    original_text = "this is a test"
    tokenized_text = " ".join(c for c in original_text if c != " ")
    original_words_to_tokes = {w: list(w) for w in original_text.split()}
    TOKENIZATION_RESULT = {
        "originalText": original_text,
        "tokenizedText": tokenized_text,
        "originalWordsToTokens": original_words_to_tokes,
    }
    positive_tokens = ["t", "i"]
    negative_tokens = []
    d = {"Label": "Valid", "Probability": 0.7, "PositiveWords": positive_tokens, "NegativeWords": negative_tokens}
    mocker.patch.object(phishing_mock, "explain_model_words", return_value=d, create=True)
    res = predict_phishing_words("modelName", "list", original_text, "", 0, 0, 0, 10, True)
    correct_highlighted = " ".join(
        bold(w) if any(pos_token in w for pos_token in positive_tokens) else w for w in original_text.split()
    )
    assert res["Contents"] == {
        "OriginalText": original_text,
        "Probability": 0.7,
        "NegativeWords": negative_tokens,
        "TextTokensHighlighted": correct_highlighted,
        "PositiveWords": [w for w in original_text.split() if any(pos_token in w for pos_token in positive_tokens)],
        "Label": "Valid",
    }


def test_predict_phishing_words_tokenization_by_character_hashed(mocker):
    global TOKENIZATION_RESULT

    def hash_token(t):
        return str(ord(t))

    def unhash_token(t):
        return chr(int(t))

    phishing_mock = PhishingModelMock()
    mocker.patch.object(demisto, "executeCommand", side_effect=executeCommand)
    mocker.patch.object(demisto, "args", return_value={"topWordsLimit": 10, "hashSeed": 10})
    mocker.patch("demisto_ml.phishing_model_loads_handler", return_value=phishing_mock, create=True)
    mocker.patch.object(phishing_mock, "filter_model_words", return_value=("text", 2), create=True)
    mocker.patch.object(demisto, "incidents", return_value=[{"isPlayground": True}])
    original_text = "this is a test"
    tokenized_text = " ".join(c for c in original_text if c != " ")
    original_words_to_tokes = {w: list(w) for w in original_text.split()}
    TOKENIZATION_RESULT = {
        "originalText": original_text,
        "tokenizedText": tokenized_text,
        "originalWordsToTokens": original_words_to_tokes,
        "hashedTokenizedText": "".join(hash_token(t) if t != " " else t for t in tokenized_text),
        "wordsToHashedTokens": {w: [hash_token(t) for t in token_list] for w, token_list in original_words_to_tokes.items()},
    }
    positive_tokens = [hash_token("t"), hash_token("i")]
    negative_tokens = []
    d = {"Label": "Valid", "Probability": 0.7, "PositiveWords": positive_tokens, "NegativeWords": negative_tokens}
    mocker.patch.object(phishing_mock, "explain_model_words", return_value=d, create=True)
    res = predict_phishing_words("modelName", "list", original_text, "", 0, 0, 0, 10, True)
    correct_highlighted = " ".join(
        bold(w) if any(unhash_token(pos_token) in w for pos_token in positive_tokens) else w for w in original_text.split()
    )
    assert res["Contents"] == {
        "OriginalText": original_text,
        "Probability": 0.7,
        "NegativeWords": negative_tokens,
        "TextTokensHighlighted": correct_highlighted,
        "PositiveWords": [w for w in original_text.split() if any(unhash_token(pos_token) in w for pos_token in positive_tokens)],
        "Label": "Valid",
    }


def test_main(mocker):
    global TOKENIZATION_RESULT
    phishing_mock = PhishingModelMock()
    d = {"Label": "Valid", "Probability": 0.7, "PositiveWords": ["word1"], "NegativeWords": ["word2"]}
    args = {
        "modelName": "modelName",
        "modelStoreType": "list",
        "emailSubject": "word1",
        "emailBody": "word2 word3",
        "minTextLength": "0",
        "labelProbabilityThreshold": "0",
        "wordThreshold": "0",
        "topWordsLimit": "10",
        "returnError": "true",
    }
    mocker.patch.object(demisto, "args", return_value=args)
    mocker.patch.object(demisto, "incidents", return_value=[{"isPlayground": True}])
    mocker.patch.object(demisto, "executeCommand", side_effect=executeCommand)
    mocker.patch("demisto_ml.phishing_model_loads_handler", return_value=phishing_mock, create=True)
    mocker.patch.object(phishing_mock, "filter_model_words", return_value=("text", 2), create=True)
    mocker.patch.object(phishing_mock, "explain_model_words", return_value=d, create=True)

    TOKENIZATION_RESULT = {
        "originalText": "{} {}".format(args["emailSubject"], args["emailBody"]),
        "tokenizedText": "{} {}".format(args["emailSubject"], args["emailBody"]),
        "originalWordsToTokens": {"word1": ["word1"], "word2": ["word2"], "word3": ["word3"]},
    }

    res = main()
    correct_res = {
        "OriginalText": concatenate_subject_body(args["emailSubject"], args["emailBody"]),
        "Probability": 0.7,
        "NegativeWords": ["word2"],
        "TextTokensHighlighted": concatenate_subject_body(bold(args["emailSubject"]), args["emailBody"]),
        "PositiveWords": ["word1"],
        "Label": "Valid",
    }
    assert res["Contents"] == correct_res

    args["emailBodyHTML"] = args.pop("emailBody")
    TOKENIZATION_RESULT = {
        "originalText": concatenate_subject_body(args["emailSubject"], args["emailBodyHTML"]),
        "tokenizedText": concatenate_subject_body(args["emailSubject"], args["emailBodyHTML"]),
        "originalWordsToTokens": {"word1": ["word1"], "word2": ["word2"], "word3": ["word3"]},
    }
    main()
    assert res["Contents"] == correct_res


def test_no_positive_words(mocker):
    # make sure that if no positive words were found, TextTokensHighlighted output is equivalent to original text
    global TOKENIZATION_RESULT
    phishing_mock = PhishingModelMock()
    d = {"Label": "Valid", "Probability": 0.7, "PositiveWords": [], "NegativeWords": ["word2"]}
    args = {
        "modelName": "modelName",
        "modelStoreType": "list",
        "emailSubject": "word1",
        "emailBody": "word2 word3",
        "minTextLength": "0",
        "labelProbabilityThreshold": "0",
        "wordThreshold": "0",
        "topWordsLimit": "10",
        "returnError": "true",
    }
    mocker.patch.object(demisto, "args", return_value=args)
    mocker.patch.object(demisto, "executeCommand", side_effect=executeCommand)
    phishing_mock = PhishingModelMock(("text", 2))
    mocker.patch("demisto_ml.phishing_model_loads_handler", return_value=phishing_mock, create=True)
    mocker.patch.object(demisto, "incidents", return_value=[{"isPlayground": True}])
    mocker.patch.object(phishing_mock, "filter_model_words", return_value=("text", 2), create=True)
    mocker.patch.object(phishing_mock, "explain_model_words", return_value=d, create=True)

    TOKENIZATION_RESULT = {
        "originalText": "{} {}".format(args["emailSubject"], args["emailBody"]),
        "tokenizedText": "{} {}".format(args["emailSubject"], args["emailBody"]),
        "originalWordsToTokens": {"word1": ["word1"], "word2": ["word2"], "word3": ["word3"]},
    }

    res = main()
    assert res["Contents"]["TextTokensHighlighted"] == TOKENIZATION_RESULT["originalText"]