ActiveMQ

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

import stomp


""" GLOBAL VARS """

HOSTNAME = demisto.params()["hostname"]
PORT = int(demisto.params().get("port") or 61613)
USERNAME = demisto.params().get("credentials", {}).get("identifier")
PASSWORD = demisto.params().get("credentials", {}).get("password")
CLIENT_CERT = demisto.params().get("client_cert")
CLIENT_KEY = demisto.params().get("client_key")
ROOT_CA = demisto.params().get("root_ca")


class MsgListener(stomp.ConnectionListener):
    def __init__(self):
        self.result_arr = []
        self.msg_ids = []

    def on_error(self, frame):
        demisto.results(f'received an error "{frame}"')

    def on_message(self, frame):
        self.result_arr.append(frame.body)
        self.msg_ids.append(frame.headers["message-id"])


""" HELPER FUNCTIONS """


def create_connection(client_cert, client_key, root_ca):
    client_path = None
    client_key_path = None
    if client_cert:
        client_path = "client.cert"
        with open(client_path, "wb") as file:
            file.write(client_cert)
            client_path = os.path.abspath(client_path)
    if client_key:
        client_key_path = "client_key.key"
        with open(client_key_path, "wb") as file:
            file.write(client_key)
    if root_ca:
        root_ca_path = "root_ca.key"
        with open(root_ca_path, "wb") as file:
            file.write(root_ca)

    conn = stomp.Connection([(HOSTNAME, PORT)])
    if client_cert or client_key or root_ca:
        demisto.debug("adding ssl certificate")
        conn.set_ssl(for_hosts=[(HOSTNAME, PORT)], key_file=client_key, cert_file=client_path)

    return conn


def connect(conn, client_id=None):
    if USERNAME or PASSWORD:
        if client_id and len(client_id) > 0:
            conn.connect(USERNAME, PASSWORD, wait=True, headers={"client-id": client_id})
        else:
            conn.connect(USERNAME, PASSWORD, wait=True)

    elif CLIENT_KEY or CLIENT_CERT or ROOT_CA:
        if client_id and len(client_id) > 0:
            conn.connect(wait=True)  # , headers = {'client-id': client_id })
        else:
            conn.connect(wait=True)
    else:
        raise ValueError("You must provide username/password or certificates")

    return conn


""" FUNCTIONS """


def send_message(conn):
    txid = conn.begin()
    dest = demisto.args()["destination"]
    body = demisto.args()["body"]

    if "headers" in demisto.args():
        try:
            headers_demisto = json.loads(demisto.args()["headers"])
        except Exception as e:
            demisto.error(f'Failed to parse "headers". Error: {e}')
            raise ValueError('Failed to parse "headers" argument to JSON. "headers"={}'.format(demisto.args()["headers"]))

        conn.send(dest, body, transaction=txid, headers=headers_demisto)
    else:
        conn.send(dest, body, transaction=txid)

    conn.commit(txid)
    demisto.results("Message sent to ActiveMQ destination: " + dest + " with transaction ID: " + txid)


def subscribe(client, conn, subscription_id, topic_name, queue_name):
    if not queue_name and not topic_name:
        raise ValueError("To subscribe you must provide either queue-name or topic-name")
    elif queue_name and topic_name:
        raise ValueError("Can't provide both queue-name and topic-name.")

    listener = MsgListener()
    if client and len(client) > 0:
        conn.set_listener("Demisto", listener)

    # ack='client-individual', headers={'activemq.subscriptionName': client})
    if queue_name:
        conn.subscribe("/queue/" + queue_name, subscription_id, ack="client-individual")
    elif topic_name:
        conn.subscribe(
            "/topic/" + topic_name, subscription_id, ack="client-individual", headers={"activemq.subscriptionName": client}
        )

    time.sleep(1)

    for msg in listener.result_arr:
        demisto.results(msg)

    for msg_id in listener.msg_ids:
        conn.ack(msg_id, subscription_id)


def fetch_incidents(client, conn, subscription_id, queue_name, topic_name):
    if not queue_name and not topic_name:
        raise ValueError("To fetch incidents you must provide either Queue Name or Topic Name")
    if queue_name and topic_name:
        raise ValueError("Can't provide both Queue Name and Topic name.")

    # conn = stomp.Connection(heartbeats=(4000, 4000))
    listener = MsgListener()
    if client and len(client) > 0:
        conn.set_listener("Demisto", listener)

    if queue_name:
        conn.subscribe("/queue/" + queue_name, subscription_id, ack="client-individual")
    else:
        conn.subscribe(
            "/topic/" + topic_name, subscription_id, ack="client-individual", headers={"activemq.subscriptionName": client}
        )

    incidents = []
    time.sleep(10)
    for i in range(len(listener.result_arr)):
        msg = listener.result_arr[i]
        msg_id = listener.msg_ids[i]
        incidents.append({"Name": "ActiveMQ incident:" + msg_id, "rawJSON": msg, "details": msg})
    demisto.incidents(incidents)

    for msg_id in listener.msg_ids:
        conn.ack(msg_id, subscription_id)


def main():
    client = demisto.params().get("client-id", "Demisto")

    conn = create_connection(client_cert=CLIENT_CERT, client_key=CLIENT_KEY, root_ca=ROOT_CA)

    LOG(f"command is {demisto.command()}")

    try:
        if demisto.command() == "test-module":
            # Test connectivity
            if demisto.params().get("isFetch"):
                queue_name = demisto.params().get("queue_name")
                topic_name = demisto.params().get("topic-name")

                if not queue_name and not topic_name:
                    raise ValueError("To fetch incidents you must provide either Queue Name or Topic Name")
                elif queue_name and topic_name:
                    raise ValueError("Can't provide both Queue Name and Topic name.")

            connect(conn)
            demisto.results("ok")

        elif demisto.command() == "activemq-send":
            connect(conn)
            send_message(conn)

        elif demisto.command() == "activemq-subscribe":
            subscription_id = demisto.args().get("subscription-id")
            topic_name = demisto.args().get("topic-name")
            queue_name = demisto.args().get("queue-name")

            connect(conn, client)
            subscribe(client, conn, subscription_id, topic_name, queue_name)

        elif demisto.command() == "fetch-incidents":
            subscription_id = demisto.params().get("subscription-id")
            queue_name = demisto.params().get("queue_name")
            topic_name = demisto.params().get("topic-name")

            connect(conn, client)
            fetch_incidents(client, conn, subscription_id, queue_name, topic_name)

    except Exception as e:
        if demisto.command() == "fetch-incidents":
            raise

        return_error(str(e))

    finally:
        conn.disconnect()


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