AWS - Route53

from datetime import date
from http import HTTPStatus

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

# Disable insecure warnings
urllib3.disable_warnings()

SERVICE = "route53"
DEFAULT_RETRIES = 5

"""HELPER FUNCTIONS"""


class DatetimeEncoder(json.JSONEncoder):
    # pylint: disable=method-hidden
    def default(self, obj: Any) -> Any:
        if isinstance(obj, datetime):
            return obj.strftime("%Y-%m-%dT%H:%M:%S")
        elif isinstance(obj, date):
            return obj.strftime("%Y-%m-%d")
        # Let the base class default method raise the TypeError
        return json.JSONEncoder.default(self, obj)


def create_entry(title: str, data: Union[Dict[str, Any], List[Any]], outputs: Any, outputs_prefix: str) -> CommandResults:
    return CommandResults(
        entry_type=EntryType.NOTE,
        content_format=EntryFormat.JSON,
        readable_output=tableToMarkdown(title, data, removeNull=True) if data else "No result were found",
        outputs=outputs,
        outputs_prefix=outputs_prefix,
    )


def raise_error(error: Any) -> CommandResults:
    demisto.error(f"Error occurred in {SERVICE} - {error!s}")
    return CommandResults(content_format=EntryFormat.TEXT, entry_type=EntryType.ERROR, readable_output=str(error))


def create_record(
    args: Dict[Any, Any],
    aws_session: Any,
) -> CommandResults:
    try:
        change_batch: Dict[Any, Any] = {
            "Changes": [
                {
                    "Action": "CREATE",
                    "ResourceRecordSet": {
                        "Name": args.get("source"),
                        "Type": args.get("type"),
                        "TTL": arg_to_number(args.get("ttl"), "ttl", True),
                        "ResourceRecords": [{"Value": args.get("target")}],
                    },
                }
            ]
        }

        if args.get("comment"):
            change_batch["Comment"] = args.get("comment")

        kwargs = {"HostedZoneId": args.get("hostedZoneId"), "ChangeBatch": change_batch}

        response = aws_session.change_resource_record_sets(**kwargs)
        record = response["ChangeInfo"]
        data = {"Id": record["Id"], "Status": record["Status"]}

        output = json.loads(json.dumps(response["ChangeInfo"], cls=DatetimeEncoder))
        return create_entry("AWS Route53 record created", data, output, "AWS.Route53.RecordSetsChange")

    except Exception as error:
        return raise_error(error)


def delete_record(args: Dict[Any, Any], aws_session: Any) -> CommandResults:
    try:
        kwargs = {
            "HostedZoneId": args.get("hostedZoneId"),
            "ChangeBatch": {
                "Changes": [
                    {
                        "Action": "DELETE",
                        "ResourceRecordSet": {
                            "Name": args.get("source"),
                            "Type": args.get("type"),
                            "TTL": arg_to_number(args.get("ttl"), "ttl", True),
                            "ResourceRecords": [{"Value": args.get("target")}],
                        },
                    }
                ]
            },
        }

        response = aws_session.change_resource_record_sets(**kwargs)
        record = response["ChangeInfo"]
        data = {"Id": record["Id"], "Status": record["Status"]}

        output = json.loads(json.dumps(response["ChangeInfo"], cls=DatetimeEncoder))
        return create_entry("AWS Route53 record deleted", data, output, "AWS.Route53.RecordSetsChange")

    except Exception as error:
        return raise_error(error)


def upsert_record(args: Dict[Any, Any], aws_session: Any) -> CommandResults:
    try:
        change_batch: Dict[Any, Any] = {
            "Changes": [
                {
                    "Action": "UPSERT",
                    "ResourceRecordSet": {
                        "Name": args.get("source"),
                        "Type": args.get("type"),
                        "TTL": arg_to_number(args.get("ttl"), "ttl", True),
                        "ResourceRecords": [{"Value": args.get("target")}],
                    },
                }
            ]
        }

        if args.get("comment"):
            change_batch["Comment"] = args.get("comment")
        kwargs = {"HostedZoneId": args.get("hostedZoneId"), "ChangeBatch": change_batch}

        response = aws_session.change_resource_record_sets(**kwargs)
        record = response["ChangeInfo"]
        data = {"Id": record["Id"], "Status": record["Status"]}

        output = json.loads(json.dumps(response["ChangeInfo"], cls=DatetimeEncoder))
        return create_entry("AWS Route53 record Upsert", data, output, "AWS.Route53.RecordSetsChange")

    except Exception as error:
        return raise_error(error)


def list_hosted_zones(aws_session: Any) -> CommandResults:
    try:
        data = []
        response = aws_session.list_hosted_zones()
        for hosted_zone in response["HostedZones"]:
            data.append(
                {
                    "Name": hosted_zone["Name"],
                    "Id": hosted_zone["Id"],
                    "ResourceRecordSetCount": hosted_zone["ResourceRecordSetCount"],
                }
            )
        output = json.loads(json.dumps(data, cls=DatetimeEncoder))
        return create_entry("AWS Route53 Hosted Zones", data, output, "AWS.Route53.HostedZones")

    except Exception as error:
        return raise_error(error)


def list_resource_record_sets(args: Dict[Any, Any], aws_session: Any) -> CommandResults:
    try:
        kwargs = {"HostedZoneId": args.get("hostedZoneId")}
        if args.get("startRecordName"):
            kwargs.update({"StartRecordName": args.get("startRecordName")})
        if args.get("startRecordType"):
            kwargs.update({"StartRecordType": args.get("startRecordType")})
        if args.get("startRecordIdentifier"):
            kwargs.update({"StartRecordIdentifier": args.get("startRecordIdentifier")})

        data = []
        response = aws_session.list_resource_record_sets(**kwargs)
        records = response["ResourceRecordSets"]
        for record in records:
            resource_records = record.get("ResourceRecords") or []
            data.append(
                {
                    "Name": record.get("Name"),
                    "Type": record.get("Type"),
                    "TTL": record.get("TTL"),
                    "ResourceRecords": resource_records[0]["Value"] if resource_records else None,
                }
            )
        output = json.loads(json.dumps(response["ResourceRecordSets"], cls=DatetimeEncoder))
        return create_entry("AWS Route53 Record Sets", data, output, "AWS.Route53.RecordSets")

    except Exception as error:
        return raise_error(error)


def waiter_resource_record_sets_changed(args: Dict[Any, Any], aws_session: Any) -> CommandResults:
    try:
        kwargs = {"Id": args.get("id")}
        if args.get("waiterDelay") is not None:
            kwargs.update({"WaiterConfig": {"Delay": arg_to_number(args.get("waiterDelay"), "waiterDelay", True)}})
        if args.get("waiterMaxAttempts"):
            kwargs.update(
                {"WaiterConfig": {"MaxAttempts": arg_to_number(args.get("waiterMaxAttempts"), "waiterMaxAttempts", True)}}
            )

        waiter = aws_session.get_waiter("resource_record_sets_changed")
        waiter.wait(**kwargs)
        return CommandResults(entry_type=EntryType.NOTE, content_format=EntryFormat.JSON, readable_output="success")

    except Exception as error:
        return raise_error(error)


def test_dns_answer(args: Dict[Any, Any], aws_session: Any) -> CommandResults:
    try:
        kwargs = {
            "HostedZoneId": args.get("hostedZoneId"),
            "RecordName": args.get("recordName"),
            "RecordType": args.get("recordType"),
        }
        if args.get("resolverIP"):
            kwargs.update({"ResolverIP": args.get("resolverIP")})

        response = aws_session.test_dns_answer(**kwargs)
        data = {
            "Nameserver": response["Nameserver"],
            "RecordName": response["RecordName"],
            "RecordType": response["RecordType"],
            "ResponseCode": response["ResponseCode"],
            "Protocol": response["Protocol"],
        }

        return create_entry("AWS Route53 Test DNS Answer", data, response, "AWS.Route53.TestDNSAnswer")

    except Exception as error:
        return raise_error(error)


def test_module(aws_session: Any) -> CommandResults:
    try:
        response = aws_session.list_hosted_zones()
        if response["ResponseMetadata"]["HTTPStatusCode"] == HTTPStatus.OK:
            return_results("ok")

        return CommandResults(
            content_format=EntryFormat.TEXT,
            entry_type=EntryType.ERROR,
            readable_output=f"received status code {response['ResponseMetadata']['HTTPStatusCode']}",
        )

    except Exception as error:
        return raise_error(error)


def main():  # pragma: no cover
    params = demisto.params()
    command = demisto.command()
    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") or params.get("access_key")
    aws_secret_access_key = params.get("credentials", {}).get("password") or params.get("secret_key")
    verify_certificate = not params.get("insecure", True)
    timeout = params.get("timeout")
    retries = params.get("retries", DEFAULT_RETRIES)

    try:
        args = demisto.args()
        validate_params(
            True,
            aws_role_arn,
            aws_role_session_name,
            aws_access_key_id,  # noqa
            aws_secret_access_key,
        )

        aws_client = AWSClient(
            None,
            aws_role_arn,
            aws_role_session_name,  # noqa
            aws_role_session_duration,
            aws_role_policy,
            aws_access_key_id,
            aws_secret_access_key,
            verify_certificate,
            timeout,
            retries,
        )
        aws_session = aws_client.aws_session(
            service=SERVICE,
            role_arn=aws_role_arn,
            role_session_name=aws_role_session_name,
            role_session_duration=aws_role_session_duration,
        )

        demisto.info(f"Command being called is {demisto.command()}")
        if command == "test-module":
            return_results(test_module(aws_session))

        elif command == "aws-route53-create-record":
            return_results(create_record(args, aws_session))

        elif command == "aws-route53-upsert-record":
            return_results(upsert_record(args, aws_session))

        elif command == "aws-route53-delete-record":
            return_results(delete_record(args, aws_session))

        elif command == "aws-route53-list-hosted-zones":
            return_results(list_hosted_zones(aws_session))

        elif command == "aws-route53-list-resource-record-sets":
            return_results(list_resource_record_sets(args, aws_session))

        elif command == "aws-route53-waiter-resource-record-sets-changed":
            return_results(waiter_resource_record_sets_changed(args, aws_session))

        elif command == "aws-route53-test-dns-answer":
            return_results(test_dns_answer(args, aws_session))
        else:
            raise NotImplementedError(f"{command} command is not implemented.")

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


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