data "aws_caller_identity" "current" {
}

data "aws_caller_identity" "dns" {
  provider = aws.dns
}

# Create emr ec2 instance profile
resource "aws_iam_instance_profile" "emr-ec2-instance-profile" {
  name = "${var.environment}-${var.service_name}-emr-ec2-role"
  role = aws_iam_role.emr_ec2_role.name
}

# Create ec2 assume role policy document
data "aws_iam_policy_document" "emr_ec2_assume_role_policy_document" {
  statement {
    effect  = "Allow"
    actions = ["sts:AssumeRole"]

    principals {
      type = "Service"
      identifiers = [
        "ec2.amazonaws.com"
      ]
    }
  }
}

# Create ec2 iam role and attach emr_ec2_assume_role_policy_document
resource "aws_iam_role" "emr_ec2_role" {
  name               = "${var.environment}-${var.service_name}-emr-ec2-role"
  assume_role_policy = data.aws_iam_policy_document.emr_ec2_assume_role_policy_document.json
}

# EC2 service role policy document
data "aws_iam_policy_document" "ec2_service_role_policy_document" {
  # checkov:skip=CKV_AWS_111:We grant the EMR instance profile role with full EC2 permissions. AWS-managed policy for EMR is deprecated, the recommended policy from the docs is too permissive.
  # checkov:skip=CKV_AWS_107:The same as above.
  # checkov:skip=CKV_AWS_109:The same as above.

  statement {
    effect = "Allow"

    actions = [
      "ec2:*",
      "iam:ListRolePolicies",
      "iam:GetRole",
      "iam:GetRolePolicy",
      "iam:ListInstanceProfiles",
      "sdb:BatchPutAttributes",
      "sdb:Select"
    ]

    resources = [
      "*"
    ]
  }
}

# EC2 service role policy
resource "aws_iam_policy" "ec2_service_role_policy" {
  name        = "${var.environment}-${var.service_name}-ec2-service-role-policy"
  description = "${var.environment}-${var.service_name}-ec2-service-role-policy"
  policy      = data.aws_iam_policy_document.ec2_service_role_policy_document.json
}

# EC2 service role policy attachment
resource "aws_iam_role_policy_attachment" "ec2_service_role_policy_attachment" {
  role       = aws_iam_role.emr_ec2_role.name
  policy_arn = aws_iam_policy.ec2_service_role_policy.arn
}

# Create assume role policy document
data "aws_iam_policy_document" "emr_assume_service_role_policy_document" {
  statement {
    effect  = "Allow"
    actions = ["sts:AssumeRole"]

    principals {
      type = "Service"
      identifiers = [
        "elasticmapreduce.amazonaws.com",
        "application-autoscaling.amazonaws.com"
      ]
    }
  }
}

# Create assume service role and attach assume policy doc.
resource "aws_iam_role" "emr_service_role" {
  name               = "${var.environment}-${var.service_name}-emr-service-role"
  assume_role_policy = data.aws_iam_policy_document.emr_assume_service_role_policy_document.json
}

data "aws_iam_policy_document" "emr_service_role_policy_document" {
  # checkov:skip=CKV_AWS_111:We grant the EMR service role with full EC2 permissions. AWS-managed policy for EMR is deprecated, the recommended policy from the docs doesn't provide enough permissions, tag-based policy implementation failed.
  # checkov:skip=CKV_AWS_107:The same as above.
  # checkov:skip=CKV_AWS_109:The same as above.

  statement {
    effect = "Allow"

    actions = [
      "ec2:*",
      "iam:ListRolePolicies",
      "iam:GetRole",
      "iam:GetRolePolicy",
      "iam:ListInstanceProfiles",
      "sdb:BatchPutAttributes",
      "sdb:Select"
    ]

    resources = [
      "*"
    ]
  }

  statement {
    effect = "Allow"

    actions = [
      "iam:PassRole"
    ]

    resources = [
      aws_iam_role.emr_ec2_role.arn
    ]
  }
}

resource "aws_iam_policy" "emr_service_role_policy" {
  name        = "${var.environment}-${var.service_name}-emr-service-role-policy"
  description = "${var.environment}-${var.service_name}-emr-service-role-policy"
  policy      = data.aws_iam_policy_document.emr_service_role_policy_document.json
}

resource "aws_iam_role_policy_attachment" "emr_service_role_policy_attachment" {
  role       = aws_iam_role.emr_service_role.name
  policy_arn = aws_iam_policy.emr_service_role_policy.arn
}

data "aws_subnet" "the_subnet" {
  id = var.subnet_id
}

#
# Security group resources
#
resource "aws_security_group" "emr_leader" {
  vpc_id                 = data.aws_subnet.the_subnet.vpc_id
  revoke_rules_on_delete = true
  name                   = "${var.environment}-${var.service_name}-EMR-leader"

  tags = merge({
    name = "${var.environment}-${var.service_name}-EMR-leader",
  }, local.tags)

  ingress {
    description = ""
    from_port   = 0
    to_port     = 0
    protocol    = "-1"
    self        = true
  }
  ingress {
    description = ""
    from_port   = 22
    to_port     = 22
    protocol    = "tcp"
    cidr_blocks = var.leader_ingress_ssh
  }
  ingress {
    description = ""
    from_port   = 8443
    to_port     = 8443
    protocol    = "tcp"
    cidr_blocks = concat(local.tcp_8443_listener_allow_cidr_blocks[var.environment], split(",", data.aws_subnet.the_subnet.cidr_block))
  }
  ingress {
    description = ""
    from_port   = 8998
    to_port     = 8998
    protocol    = "tcp"
    cidr_blocks = concat(local.tcp_8998_listener_allow_cidr_blocks[var.environment], split(",", data.aws_subnet.the_subnet.cidr_block))
  }
  ingress {
    description = ""
    from_port   = 18080
    to_port     = 18080
    protocol    = "tcp"
    cidr_blocks = concat(local.tcp_18080_listener_allow_cidr_blocks[var.environment], split(",", data.aws_subnet.the_subnet.cidr_block))
  }
  egress {
    description = ""
    from_port   = 0
    to_port     = 0
    protocol    = "-1"
    cidr_blocks = ["0.0.0.0/0"]
  }
}

resource "aws_security_group" "emr_core" {
  vpc_id                 = data.aws_subnet.the_subnet.vpc_id
  revoke_rules_on_delete = true
  name                   = "${var.environment}-${var.service_name}-EMR-core"

  tags = merge({
    name = "${var.environment}-${var.service_name}-EMR-core",
  }, local.tags)

  dynamic "ingress" {
    # checkov:skip=CKV_AWS_24:We allow SSH access only if a list of source CIDR blocks is provided.
    for_each = length(var.core_ingress_ssh) > 0 ? ["ssh"] : []
    content {
      description = ""
      from_port   = 22
      to_port     = 22
      protocol    = "tcp"
      cidr_blocks = var.core_ingress_ssh
    }
  }

  ingress {
    description = ""
    from_port   = 0
    to_port     = 0
    protocol    = "-1"
    self        = true
  }

  egress {
    description = ""
    from_port   = 0
    to_port     = 0
    protocol    = "-1"
    cidr_blocks = ["0.0.0.0/0"]
  }
}

resource "aws_security_group" "emr_cluster" {
  vpc_id                 = data.aws_subnet.the_subnet.vpc_id
  revoke_rules_on_delete = true
  name                   = "${var.environment}-${var.service_name}-EMR-cluster"

  ingress {
    description     = ""
    from_port       = 0
    to_port         = 0
    protocol        = "-1"
    security_groups = [aws_security_group.emr_leader.id]
  }
  ingress {
    description     = ""
    from_port       = 0
    to_port         = 0
    protocol        = "-1"
    security_groups = [aws_security_group.emr_core.id]
  }
  egress {
    description = ""
    from_port   = 0
    to_port     = 0
    protocol    = "-1"
    cidr_blocks = ["0.0.0.0/0"]
  }
}

resource "aws_security_group" "emr_service" {
  vpc_id                 = data.aws_subnet.the_subnet.vpc_id
  revoke_rules_on_delete = true
  name                   = "${var.environment}-${var.service_name}-EMR-Service"

  tags = merge({
    name = "${var.environment}-${var.service_name}-EMR-Service",
  }, local.tags)

  egress {
    description     = ""
    from_port       = 8443
    to_port         = 8443
    protocol        = "tcp"
    security_groups = [aws_security_group.emr_leader.id]
  }
  egress {
    description     = ""
    from_port       = 8443
    to_port         = 8443
    protocol        = "tcp"
    security_groups = [aws_security_group.emr_core.id]
  }
}

#
# EMR resources
#
resource "aws_emr_cluster" "cluster" {
  name                              = "${var.environment}-${var.service_name}-cluster"
  release_label                     = var.release_label
  applications                      = var.applications
  configurations                    = var.configurations
  keep_job_flow_alive_when_no_steps = var.job_flow_keep_alive

  ec2_attributes {
    key_name                          = var.key_name
    subnet_id                         = var.subnet_id
    emr_managed_master_security_group = aws_security_group.emr_leader.id
    additional_master_security_groups = aws_security_group.emr_cluster.id
    emr_managed_slave_security_group  = aws_security_group.emr_core.id
    additional_slave_security_groups  = aws_security_group.emr_cluster.id
    service_access_security_group     = aws_security_group.emr_service.id
    instance_profile                  = aws_iam_instance_profile.emr-ec2-instance-profile.arn
  }

  master_instance_group {
    name           = var.leader_instance_group_name
    instance_type  = var.leader_instance_group_instance_type
    instance_count = var.leader_instance_group_instance_count

    ebs_config {
      size                 = var.leader_instance_group_ebs_size
      type                 = var.leader_instance_group_ebs_type
      iops                 = var.leader_instance_group_ebs_iops
      volumes_per_instance = var.leader_instance_group_ebs_volumes_per_instance
    }
  }

  core_instance_group {
    name           = var.core_instance_group_name
    instance_type  = var.core_instance_group_instance_type
    instance_count = var.core_instance_group_instance_count

    ebs_config {
      size                 = var.leader_instance_group_ebs_size
      type                 = var.leader_instance_group_ebs_type
      iops                 = var.leader_instance_group_ebs_iops
      volumes_per_instance = var.leader_instance_group_ebs_volumes_per_instance
    }
  }

  bootstrap_action {
    path = var.bootstrap_uri
    name = var.bootstrap_name
    args = var.bootstrap_args
  }

  log_uri      = var.log_uri
  service_role = aws_iam_role.emr_service_role.arn

  tags = merge({
    Name = "${var.environment}-${var.service_name}-emr",
  }, local.tags)

  security_configuration = var.kms_key_enabled ? aws_emr_security_configuration.security_configuration[0].name : ""
}

resource "aws_route53_record" "emr_cluster_route53_record" {
  provider = aws.dns
  count    = var.route53_record_creation_enabled ? 1 : 0

  name    = "${var.environment}-${var.service_name}-emr"
  zone_id = local.route53_zone_id
  ttl     = var.route53_record_ttl
  type    = local.leader_node_hostname_is_internal ? "A" : "CNAME"
  records = local.leader_node_hostname_is_internal ? [local.leader_node_ip] : [aws_emr_cluster.cluster.master_public_dns]
}

resource "aws_security_group_rule" "managed_master_service_access_ingress" {
  description              = "Allow ingress traffic from EmrManagedLeaderSecurityGroup"
  type                     = "ingress"
  from_port                = 9443
  to_port                  = 9443
  protocol                 = "tcp"
  source_security_group_id = aws_security_group.emr_leader.id
  security_group_id        = aws_security_group.emr_service.id
}

# Attach additional existing managed IAM policies to emr service role, if enabled
resource "aws_iam_role_policy_attachment" "emr_service_role_existing_iam_policy_attachment" {
  count      = length(var.emr_service_role_iam_managed_policy_attachments)
  role       = aws_iam_role.emr_service_role.name
  policy_arn = var.emr_service_role_iam_managed_policy_attachments[count.index]
}

# Attach additional existing managed IAM policies to ec2 role, if enabled
resource "aws_iam_role_policy_attachment" "emr_ec2_service_role_existing_iam_policy_attachment" {
  count      = length(var.emr_ec2_role_iam_managed_policy_attachments)
  role       = aws_iam_role.emr_ec2_role.name
  policy_arn = var.emr_ec2_role_iam_managed_policy_attachments[count.index]
}

# Create KMS key and policies for encrypt EMR EBS volume
resource "aws_kms_key" "emr_kms_key" {
  count = var.kms_key_enabled ? 1 : 0

  description         = "${var.environment}-${var.service_name}-EMR-Encryption"
  enable_key_rotation = true
  tags                = local.tags
}

# Creating the key alias
resource "aws_kms_alias" "emr_kms_alias" {
  count = var.kms_key_enabled ? 1 : 0

  name          = "alias/${var.environment}-${var.service_name}-emr-kms-key"
  target_key_id = element(aws_kms_key.emr_kms_key[*].key_id, count.index)
}

# IAM policy document
data "aws_iam_policy_document" "kms_emr_policy_document" {
  statement {
    effect = "Allow"
    actions = [
      "kms:Encrypt",
      "kms:Decrypt",
      "kms:ReEncrypt*",
      "kms:CreateGrant",
      "kms:GenerateDataKey*",
      "kms:DescribeKey",
    ]
    resources = aws_kms_key.emr_kms_key[*].arn
  }
}

# Create KMS policy
resource "aws_iam_policy" "emr_kms_policy" {
  count = var.kms_key_enabled ? 1 : 0

  name   = "KMS-${var.environment}-${var.service_name}-EMR-Encryption-policy"
  policy = data.aws_iam_policy_document.kms_emr_policy_document.json
  tags   = local.tags
}

# Policy attachments to EMR roles
resource "aws_iam_role_policy_attachment" "emr_ec2_role_kms_policy_attachment" {
  count = var.kms_key_enabled ? 1 : 0

  role       = aws_iam_role.emr_ec2_role.name
  policy_arn = element(aws_iam_policy.emr_kms_policy[*].arn, count.index)
}

resource "aws_iam_role_policy_attachment" "emr_service_role_kms_policy_attachment" {
  count = var.kms_key_enabled ? 1 : 0

  role       = aws_iam_role.emr_service_role.name
  policy_arn = element(aws_iam_policy.emr_kms_policy[*].arn, count.index)
}

# Creating EMR security configuaration
resource "aws_emr_security_configuration" "security_configuration" {
  count = var.kms_key_enabled ? 1 : 0

  name = "${var.environment}-${var.service_name}-emr-security-configuration"

  configuration = jsonencode({
    EncryptionConfiguration = {
      EnableInTransitEncryption = false,
      EnableAtRestEncryption    = true,

      AtRestEncryptionConfiguration = {
        S3EncryptionConfiguration = {
          EncryptionMode = "SSE-KMS",
          AwsKmsKey      = element(aws_kms_key.emr_kms_key[*].arn, count.index)
        },

        LocalDiskEncryptionConfiguration = {
          EnableEbsEncryption       = true,
          EncryptionKeyProviderType = "AwsKms",
          AwsKmsKey                 = element(aws_kms_key.emr_kms_key[*].arn, count.index)
        },
      }
    }
  })
}