AWS - SNS

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


def create_entry(title, data, ec):  # pragme no cover
    return {
        "ContentsFormat": formats["json"],
        "Type": entryTypes["note"],
        "Contents": data,
        "ReadableContentsFormat": formats["markdown"],
        "HumanReadable": tableToMarkdown(title, data) if data else "No result were found",
        "EntryContext": ec,
    }


def raise_error(error):  # pragma no cover
    return {"Type": entryTypes["error"], "ContentsFormat": formats["text"], "Contents": str(error)}


def create_subscription(args, client):
    try:
        attributes = {}
        kwargs = {"TopicArn": args.get("topicArn"), "Protocol": args.get("protocol")}
        if args.get("endpoint") is not None:
            kwargs.update({"Endpoint": args.get("endpoint")})
        if args.get("returnSubscriptionArn") is not None:
            kwargs.update({"ReturnSubscriptionArn": bool(args.get("returnSubscriptionArn"))})

        if args.get("deliveryPolicy") is not None:
            attributes.update({"DeliveryPolicy": args.get("deliveryPolicy")})
        if args.get("filterPolicy") is not None:
            attributes.update({"FilterPolicy": args.get("filterPolicy")})
        if args.get("rawMessageDelivery") is not None:
            attributes.update({"RawMessageDelivery": args.get("rawMessageDelivery")})
        if args.get("redrivePolicy") is not None:
            attributes.update({"RedrivePolicy": args.get("RedrivePolicy")})
        if args.get("subscriptionRoleArn") is not None:
            attributes.update({"SubscriptionRoleArn": args.get("subscriptionRoleArn")})
        if attributes:
            kwargs.update({"Attributes": attributes})

        response = client.subscribe(**kwargs)
        data = {"SubscriptionArn": response["SubscriptionArn"]}

        ec = {"AWS.SNS.Subscriptions": data}
        return create_entry("AWS SNS Subscriptions", data, ec)

    except Exception as e:
        return raise_error(e)


def list_topics(args, client):
    try:
        data = []
        kwargs = {}
        if args.get("nextToken") is not None:
            kwargs.update({"NextToken": args.get("nextToken")})
        response = client.list_topics(**kwargs)
        for topic in response["Topics"]:
            data.append({"TopicArn": topic})

        ec = {"AWS.SNS.Topics": data}
        return create_entry("AWS SNS Topics", data, ec)

    except Exception as e:
        return raise_error(e)


def list_subscriptions_by_topic(args, client):
    try:
        data = []
        kwargs = {}
        if args.get("topicArn") is not None:
            kwargs.update({"TopicArn": args.get("topicArn")})
        if args.get("nextToken") is not None:
            kwargs.update({"NextToken": args.get("nextToken")})
        response = client.list_subscriptions_by_topic(**kwargs)
        for subscription in response["Subscriptions"]:
            data.append({"SubscriptionArn": subscription["SubscriptionArn"]})

        ec = {"AWS.SNS.Subscriptions": data}
        return create_entry("AWS SNS Subscriptions", data, ec)

    except Exception as e:
        return raise_error(e)


def send_message(args, client):
    try:
        data = []
        kwargs = {"Message": args.get("message")}

        if args.get("topicArn") is not None:
            kwargs.update({"TopicArn": args.get("topicArn")})
        if args.get("targetArn") is not None:
            kwargs.update({"TargetArn": args.get("targetArn")})
        if args.get("phoneNumber") is not None:
            kwargs.update({"PhoneNumber": args.get("phoneNumber")})
        if args.get("subject") is not None:
            kwargs.update({"Subject": args.get("subject")})
        if args.get("messageStructure") is not None:
            kwargs.update({"MessageStructure": args.get("messageStructure")})
        if args.get("messageDeduplicationId") is not None:
            kwargs.update({"MessageDeduplicationId": args.get("messageDeduplicationId")})
        if args.get("messageGroupId") is not None:
            kwargs.update({"MessageGroupId": args.get("messageGroupId")})

        response = client.publish(**kwargs)
        data.append({"MessageId": response["MessageId"]})
        ec = {"AWS.SNS.SentMessages": data}
        return create_entry("AWS SNS sent messages", data, ec)

    except Exception as e:
        return raise_error(e)


def create_topic(args, client):
    try:
        attributes = {}
        kwargs = {"Name": args.get("topicName")}
        if args.get("deliveryPolicy") is not None:
            attributes.update({"DeliveryPolicy": args.get("deliveryPolicy")})
        if args.get("displayName") is not None:
            attributes.update({"DisplayName": args.get("displayName")})
        if args.get("fifoTopic") is not None:
            attributes.update({"FifoTopic": bool(args.get("fifoTopic"))})
        if args.get("policy") is not None:
            attributes.update({"policy": args.get("Policy")})
        if args.get("kmsMasterKeyId") is not None:
            attributes.update({"KmsMasterKeyId": args.get("kmsMasterKeyId")})
        if args.get("contentBasedDeduplication") is not None:
            attributes.update({"ContentBasedDeduplication": args.get("contentBasedDeduplication")})
        if attributes:
            kwargs.update({"Attributes": attributes})

        response = client.create_topic(**kwargs)
        data = {"ARN": response["TopicArn"]}
        ec = {"AWS.SNS.Topic": data}
        return create_entry("AWS SNS Topic", data, ec)

    except Exception as e:
        return raise_error(e)


def delete_topic(args, client):
    try:
        response = client.delete_topic(TopicArn=args.get("topicArn"))
        if response["ResponseMetadata"]["HTTPStatusCode"] == 200:
            return "The Topic has been deleted"

    except Exception as e:
        return raise_error(e)


def test_function(aws_client):  # pragma no cover
    try:
        client = aws_client.aws_session(service="sns")
        response = client.list_topics()
        if response["ResponseMetadata"]["HTTPStatusCode"] == 200:
            return "ok"
    except Exception as e:
        return raise_error(e)


def main():  # pragma no cover
    params = demisto.params()
    aws_default_region = params.get("defaultRegion")
    aws_role_arn = params.get("roleArn")
    aws_role_session_name = params.get("roleSessionName")
    aws_role_session_duration = params.get("sessionDuration")
    aws_role_policy = None
    aws_access_key_id = params.get("credentials", {}).get("identifier", "")
    aws_secret_access_key = params.get("credentials", {}).get("password", "")
    verify_certificate = not params.get("insecure", False)
    timeout = params.get("timeout")
    retries = params.get("retries") or 5

    commands = {
        "aws-sns-create-subscription": create_subscription,
        "aws-sns-list-topics": list_topics,
        "aws-sns-list-subscriptions-by-topic": list_subscriptions_by_topic,
        "aws-sns-send-message": send_message,
        "aws-sns-create-topic": create_topic,
        "aws-sns-delete-topic": delete_topic,
    }

    try:
        validate_params(aws_default_region, aws_role_arn, aws_role_session_name, aws_access_key_id, aws_secret_access_key)
        aws_client = AWSClient(
            aws_default_region,
            aws_role_arn,
            aws_role_session_name,
            aws_role_session_duration,
            aws_role_policy,
            aws_access_key_id,
            aws_secret_access_key,
            verify_certificate,
            timeout,
            retries,
        )
        command = demisto.command()
        args = demisto.args()
        demisto.debug(f"Command being called is {command}")
        if command == "test-module":
            return_results(test_function(aws_client))
        elif command in commands:
            client = aws_client.aws_session(
                service="sns",
                region=args.get("region"),
                role_arn=args.get("roleArn"),
                role_session_name=args.get("roleSessionName"),
                role_session_duration=args.get("roleSessionDuration"),
            )
            return_results(commands[command](args, client))
        else:
            raise NotImplementedError(f"{command} is not an existing AWS-SNS command")

    except Exception as e:
        return_error(f"Failed to execute {demisto.command()} command.\nError:\n{e!s}")


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