ATDDetonate

Detonate File or URL through McAfee ATD.

Type

python

Pack

McAfee_Advanced_Threat_Defense

Source

from time import sleep

import demistomock as demisto  # noqa: F401
from CommonServerPython import *  # noqa: F401


def main():
    dArgs = demisto.args()

    timeout = 960
    if "timeout" in demisto.args():
        timeout = int(demisto.args()["timeout"])
    interval = 10
    if "interval" in demisto.args():
        interval = int(demisto.args()["interval"])

    taskID = ""

    # Upload file/url and get taskID
    resp = demisto.executeCommand("atd-file-upload", dArgs)
    if isError(resp[0]):
        demisto.results(resp)
        sys.exit(0)
    else:
        upload_res = demisto.get(resp[0], "Contents.results")
        if upload_res and isinstance(upload_res, list):
            taskID = demisto.get(upload_res[0], "taskId")
        else:
            demisto.results(
                {"Type": entryTypes["error"], "ContentsFormat": formats["text"], "Contents": "Coudn't extract TaskID from upload"}
            )
            sys.exit(0)

    if taskID == "-1":
        demisto.results({"Type": entryTypes["error"], "ContentsFormat": formats["text"], "Contents": "File type not supported"})
        sys.exit(0)

    # Poll stage ############
    status = None
    istate = None
    sec = 0

    atdDone = False
    while sec < timeout and not atdDone:
        # Get status
        resp = demisto.executeCommand("atd-check-status", {"taskId": taskID})
        if isError(resp[0]):
            demisto.results(resp)
            sys.exit(0)

        status = demisto.get(resp[0], "Contents.results.status")
        istate = demisto.get(resp[0], "Contents.results.istate")

        # find status
        if istate and int(istate) in [1, 2]:
            atdDone = True
        # continue loop
        else:
            sec += interval
            sleep(interval)  # pylint: disable=sleep-exists
    # Get results ############
    if not atdDone:
        demisto.results(
            {
                "Type": entryTypes["error"],
                "ContentsFormat": formats["text"],
                "Contents": f"Could not retrieve results from ATD (may be due to timeout). last status = {status}",
            }
        )
        sys.exit(0)
    if istate and int(istate) in [1, 2]:
        reportType = "json"
        if demisto.get(demisto.args(), "reportType"):
            reportType = demisto.args()["reportType"]
        demisto.results(demisto.executeCommand("atd-get-report", {"taskId": taskID, "type": reportType}))
    else:
        demisto.results(
            {
                "Type": entryTypes["error"],
                "ContentsFormat": formats["text"],
                "Contents": f"ATD: Failed to detonate source, exit status = {status}",
            }
        )


if __name__ == "__builtin__" or __name__ == "builtins":
    main()

README