McAfee DXL

import tempfile
import demistomock as demisto  # noqa: F401
from CommonServerPython import *  # noqa: F401
from dxlclient.broker import Broker
from dxlclient.client import DxlClient
from dxlclient.client_config import DxlClientConfig
from dxlclient.message import Event

from CommonServerUserPython import *

INTEGRATION_NAME = "McAfee DXL"
CONNECT_RETRIES = 1
RECONNECT_DELAY = 1
RECONNECT_DELAY_MAX = 10


class EventSender:
    TRUST_LEVEL = {
        "NOT_SET": "0",
        "KNOWN_MALICIOUS": "1",
        "MOST_LIKELY_MALICIOUS": "15",
        "MIGHT_BE_MALICIOUS": "30",
        "UNKNOWN": "50",
        "MIGHT_BE_TRUSTED": "70",
        "MOST_LIKELY_TRUSTED": "85",
        "KNOWN_TRUSTED": "99",
        "KNOWN_TRUSTED_INSTALLER": "100",
    }
    broker_ca_bundle = tempfile.NamedTemporaryFile().name
    cert_file = tempfile.NamedTemporaryFile().name
    private_key = tempfile.NamedTemporaryFile().name

    def __init__(self, params: dict):
        with open(self.broker_ca_bundle, "w") as text_file:
            text_file.write(params["broker_ca_bundle"])
        with open(self.cert_file, "w") as text_file:
            text_file.write(params["cert_file"])
        with open(self.private_key, "w") as text_file:
            text_file.write(params["private_key"])

        if "broker_urls" in params:
            self.broker_urls = params["broker_urls"].split(",")
        self.push_ip_topic = params.get("push_ip_topic")
        self.push_url_topic = params.get("push_url_topic")
        self.push_domain_topic = params.get("push_domain_topic")
        self.push_hash_topic = params.get("push_hash_topic")
        self.client = DxlClient(self.get_client_config())
        self.client.connect()

    def __del__(self):
        self.client.disconnect()

    def push_ip(self, ip, trust_level, topic):
        if not is_ip_valid(ip):
            raise ValueError(f"argument ip {ip} is not a valid IP")

        trust_level_key = self.TRUST_LEVEL[trust_level]
        if topic:
            self.push_ip_topic = topic

        self.send_event(self.push_ip_topic, f"ip:{ip};trust_level:{trust_level_key}")
        return f"Successfully pushed ip {ip} with trust level {trust_level}"

    def push_url(self, url, trust_level, topic):
        trust_level_key = self.TRUST_LEVEL[trust_level]
        if topic:
            self.push_url_topic = topic

        self.send_event(self.push_url_topic, f"url:{url};trust_level:{trust_level_key}")
        return f"Successfully pushed url {url} with trust level {trust_level}"

    def push_domain(self, domain, trust_level, topic):
        trust_level_key = self.TRUST_LEVEL[trust_level]
        if topic:
            self.push_domain_topic = topic

        self.send_event(self.push_domain_topic, f"domain:{domain};trust_level:{trust_level_key}")
        return f"Successfully pushed domain {domain} with trust level {trust_level}"

    def push_hash(self, hash_obj, trust_level, topic):
        trust_level_key = self.TRUST_LEVEL[trust_level]
        if topic:
            self.push_ip_topic = topic

        self.send_event(self.push_hash_topic, f"hash:{hash_obj};trust_level:{trust_level_key}")
        return f"Successfully pushed hash {hash_obj} with trust level {trust_level}"

    def get_client_config(self):
        config = DxlClientConfig(
            broker_ca_bundle=self.broker_ca_bundle,
            cert_file=self.cert_file,
            private_key=self.private_key,
            brokers=[Broker.parse(url) for url in self.broker_urls],
        )
        config.connect_retries = CONNECT_RETRIES
        config.reconnect_delay = RECONNECT_DELAY
        config.reconnect_delay_max = RECONNECT_DELAY_MAX
        return config

    def send_event(self, topic, payload):
        if not topic:
            raise Exception(f"Error in {demisto.command()} topic field is required")

        event = Event(topic)
        event.payload = str(payload).encode()
        self.client.send_event(event)

    def send_event_wrapper(self, topic, payload):
        self.send_event(topic, payload)
        return "Successfully sent event"


def validate_certificates_format():
    if "-----BEGIN PRIVATE KEY-----" not in demisto.params()["private_key"]:
        return_error("The private key content seems to be incorrect as it doesn't start with -----BEGIN PRIVATE KEY-----")
    if "-----END PRIVATE KEY-----" not in demisto.params()["private_key"]:
        return_error("The private key content seems to be incorrect as it doesn't end with -----END PRIVATE KEY-----")
    if "-----BEGIN CERTIFICATE-----" not in demisto.params()["cert_file"]:
        return_error(
            "The client certificates content seem to be incorrect as they don't start with '-----BEGIN CERTIFICATE-----'"
        )
    if "-----END CERTIFICATE-----" not in demisto.params()["cert_file"]:
        return_error("The client certificates content seem to be incorrect as it doesn't end with -----END CERTIFICATE-----")
    if not demisto.params()["broker_ca_bundle"].lstrip(" ").startswith("-----BEGIN CERTIFICATE-----"):
        return_error("The broker certificate seem to be incorrect as they don't start with '-----BEGIN CERTIFICATE-----'")
    if not demisto.params()["broker_ca_bundle"].rstrip(" ").endswith("-----END CERTIFICATE-----"):
        return_error("The broker certificate seem to be incorrect as they don't end with '-----END CERTIFICATE-----'")


def main():
    args = demisto.args()
    command = demisto.command()
    try:
        event_sender = EventSender(demisto.params())
        result = ""
        if command == "test-module":
            event_sender.send_event("TEST", "test")
            result = "ok"
        elif command == "dxl-send-event":
            result = event_sender.send_event_wrapper(args.get("topic"), args.get("payload"))
        elif command == "dxl-push-ip":
            result = event_sender.push_ip(args.get("ip"), args.get("trust_level"), args.get("topic"))
        elif command == "dxl-push-url":
            result = event_sender.push_url(args.get("url"), args.get("trust_level"), args.get("topic"))
        elif command == "dxl-push-domain":
            result = event_sender.push_domain(args.get("domain"), args.get("trust_level"), args.get("topic"))
        elif command == "dxl-push-hash":
            result = event_sender.push_hash(args.get("hash"), args.get("trust_level"), args.get("topic"))
        else:
            raise Exception(f"{demisto.command()} is not a command")

        return_outputs(result)
    except Exception as error:
        validate_certificates_format()
        return_error(f"error in {INTEGRATION_NAME} {error!s}.", error)


if __name__ in ("__builtin__", "builtins"):
    main()