AWSTemplateFormatVersion: '2010-09-09'
Description: Amazon MSK Cluster Cost Review Remediator - review-only inventory across enabled Regions.

Parameters:
  ScheduleEnabled:
    Type: String
    Default: 'false'
    AllowedValues: ['true', 'false']
  ScheduleExpression:
    Type: String
    Default: cron(0 13 ? * MON *)
  ExcludedRegionsCsv:
    Type: String
    Default: ap-northeast-3
  IncludeClusterDetails:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
  IncludeTags:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
  IncludeNodes:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
  MaxClustersPerRegion:
    Type: Number
    Default: 100
    MinValue: 1
    MaxValue: 200
  MaxNodesPerCluster:
    Type: Number
    Default: 200
    MinValue: 1
    MaxValue: 500

Conditions:
  CreateSchedule: !Equals [!Ref ScheduleEnabled, 'true']

Resources:
  MskCostReviewRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: [lambda.amazonaws.com]
            Action: [sts:AssumeRole]
      Policies:
        - PolicyName: msk-cost-review-policy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Sid: LambdaLogging
                Effect: Allow
                Action: [logs:CreateLogGroup, logs:CreateLogStream, logs:PutLogEvents]
                Resource: '*'
              - Sid: RegionDiscovery
                Effect: Allow
                Action: [ec2:DescribeRegions]
                Resource: '*'
              - Sid: MskReadOnly
                Effect: Allow
                Action:
                  - kafka:ListClusters
                  - kafka:ListClustersV2
                  - kafka:DescribeCluster
                  - kafka:DescribeClusterV2
                  - kafka:ListNodes
                  - kafka:ListTagsForResource
                  - kafka:ListScramSecrets
                  - kafka:ListVpcConnections
                  - kafka:ListClientVpcConnections
                  - kafka:ListClusterOperations
                Resource: '*'

  MskCostReviewFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Sub msk-cost-review-${AWS::StackName}
      Runtime: python3.12
      Handler: index.lambda_handler
      Role: !GetAtt MskCostReviewRole.Arn
      Timeout: 900
      MemorySize: 256
      Environment:
        Variables:
          EXCLUDED_REGIONS_CSV: !Ref ExcludedRegionsCsv
          INCLUDE_CLUSTER_DETAILS: !Ref IncludeClusterDetails
          INCLUDE_TAGS: !Ref IncludeTags
          INCLUDE_NODES: !Ref IncludeNodes
          MAX_CLUSTERS_PER_REGION: !Ref MaxClustersPerRegion
          MAX_NODES_PER_CLUSTER: !Ref MaxNodesPerCluster
      Code:
        ZipFile: |
          import boto3
          import os
          import traceback
          from datetime import datetime, timezone
          from botocore.exceptions import ClientError, EndpointConnectionError

          def now():
              return datetime.now(timezone.utc).isoformat()

          def as_bool(value, default=False):
              if value is None:
                  return default
              return str(value).strip().lower() in ("1", "true", "yes", "y")

          def as_int(value, default):
              try:
                  return int(value)
              except Exception:
                  return default

          def csv_set(value):
              return set(x.strip() for x in str(value or "").split(",") if x.strip())

          def enabled_regions():
              ec2 = boto3.client("ec2")
              response = ec2.describe_regions(AllRegions=True)
              out = []
              for r in response.get("Regions", []):
                  status = r.get("OptInStatus")
                  if status in (None, "opt-in-not-required", "opted-in"):
                      out.append(r["RegionName"])
              return sorted(out)

          def call_with_token(method, result_key, max_items, **kwargs):
              items = []
              token = None
              while True:
                  call_args = dict(kwargs)
                  if token:
                      call_args["NextToken"] = token
                  response = method(**call_args)
                  items.extend(response.get(result_key, []) or [])
                  if len(items) >= max_items:
                      return items[:max_items], True
                  token = response.get("NextToken")
                  if not token:
                      return items, False

          def get_clusters(kafka, max_clusters):
              try:
                  return call_with_token(kafka.list_clusters_v2, "ClusterInfoList", max_clusters)
              except Exception:
                  clusters, truncated = call_with_token(kafka.list_clusters, "ClusterInfoList", max_clusters)
                  wrapped = []
                  for c in clusters:
                      wrapped.append({"Provisioned": c})
                  return wrapped, truncated

          def classify_cluster(c):
              if c.get("Provisioned"):
                  return "PROVISIONED", c.get("Provisioned")
              if c.get("Serverless"):
                  return "SERVERLESS", c.get("Serverless")
              return "UNKNOWN", c

          def get_logging_enabled(info):
              logging = info.get("LoggingInfo", {}) if isinstance(info, dict) else {}
              broker_logs = logging.get("BrokerLogs", {}) if isinstance(logging, dict) else {}
              enabled = False
              destinations = []
              for name in ("CloudWatchLogs", "Firehose", "S3"):
                  cfg = broker_logs.get(name, {}) if isinstance(broker_logs, dict) else {}
                  if isinstance(cfg, dict) and cfg.get("Enabled") is True:
                      enabled = True
                      destinations.append(name)
              return enabled, destinations

          def summarize_cluster(region, kafka, cluster_wrapper, include_tags, include_nodes, max_nodes):
              cluster_type, info = classify_cluster(cluster_wrapper)
              arn = info.get("ClusterArn")
              name = info.get("ClusterName")
              state = info.get("State")

              item = {
                  "region": region,
                  "cluster_type": cluster_type,
                  "cluster_name": name,
                  "cluster_arn": arn,
                  "state": state,
                  "creation_time": str(info.get("CreationTime")) if info.get("CreationTime") else None,
                  "current_version": info.get("CurrentVersion"),
                  "review_flags": []
              }

              if cluster_type == "PROVISIONED":
                  broker_node_group = info.get("BrokerNodeGroupInfo", {}) or {}
                  storage_info = broker_node_group.get("StorageInfo", {}) or {}
                  ebs_info = storage_info.get("EbsStorageInfo", {}) or {}
                  open_monitoring = info.get("OpenMonitoring", {}) or {}
                  prometheus = open_monitoring.get("Prometheus", {}) if isinstance(open_monitoring, dict) else {}
                  logging_enabled, log_destinations = get_logging_enabled(info)

                  item.update({
                      "kafka_version": info.get("CurrentBrokerSoftwareInfo", {}).get("KafkaVersion") if isinstance(info.get("CurrentBrokerSoftwareInfo"), dict) else None,
                      "number_of_broker_nodes": info.get("NumberOfBrokerNodes"),
                      "broker_instance_type": broker_node_group.get("InstanceType"),
                      "client_subnets_count": len(broker_node_group.get("ClientSubnets", []) or []),
                      "security_groups_count": len(broker_node_group.get("SecurityGroups", []) or []),
                      "storage_mode": info.get("StorageMode"),
                      "ebs_volume_size_gib": ebs_info.get("VolumeSize"),
                      "enhanced_monitoring": info.get("EnhancedMonitoring"),
                      "logging_enabled": logging_enabled,
                      "logging_destinations": log_destinations,
                      "open_monitoring_jmx": prometheus.get("JmxExporter", {}).get("EnabledInBroker") if isinstance(prometheus.get("JmxExporter"), dict) else None,
                      "open_monitoring_node": prometheus.get("NodeExporter", {}).get("EnabledInBroker") if isinstance(prometheus.get("NodeExporter"), dict) else None,
                      "encryption_in_transit": info.get("EncryptionInfo", {}).get("EncryptionInTransit") if isinstance(info.get("EncryptionInfo"), dict) else None,
                      "customer_action_status": info.get("CustomerActionStatus")
                  })

              elif cluster_type == "SERVERLESS":
                  vpc_configs = info.get("VpcConfigs", []) or []
                  item.update({
                      "vpc_configs_count": len(vpc_configs),
                      "client_authentication": info.get("ClientAuthentication")
                  })

              if include_tags and arn:
                  try:
                      tags = kafka.list_tags_for_resource(ResourceArn=arn).get("Tags", {})
                      item["tags"] = tags
                      item["tag_count"] = len(tags)
                  except ClientError as e:
                      item["tags_error"] = e.response.get("Error", {}).get("Code", "ClientError")

              if include_nodes and arn and cluster_type == "PROVISIONED":
                  try:
                      nodes, truncated = call_with_token(kafka.list_nodes, "NodeInfoList", max_nodes, ClusterArn=arn)
                      item["sampled_nodes"] = len(nodes)
                      item["nodes_truncated"] = truncated
                  except ClientError as e:
                      item["nodes_error"] = e.response.get("Error", {}).get("Code", "ClientError")
                  except Exception as e:
                      item["nodes_error"] = str(e)

              flags = []
              if cluster_type == "PROVISIONED":
                  flags.append("provisioned_cluster_review")
              if cluster_type == "SERVERLESS":
                  flags.append("serverless_cluster_review")
              if state and str(state).upper() not in ("ACTIVE",):
                  flags.append("cluster_state_" + str(state).lower())
              if item.get("number_of_broker_nodes", 0) and item.get("number_of_broker_nodes", 0) >= 3:
                  flags.append("multi_broker_cluster_review")
              instance_type = str(item.get("broker_instance_type") or "")
              if any(s in instance_type for s in [".large", ".xlarge", ".2xlarge", ".4xlarge", ".8xlarge", ".12xlarge", ".16xlarge"]):
                  flags.append("larger_instance_type_review")
              if item.get("ebs_volume_size_gib", 0) and item.get("ebs_volume_size_gib", 0) >= 100:
                  flags.append("large_ebs_storage_review")
              if item.get("enhanced_monitoring") and item.get("enhanced_monitoring") != "DEFAULT":
                  flags.append("enhanced_monitoring_review")
              if item.get("logging_enabled") is True:
                  flags.append("broker_logging_enabled_review")
              if item.get("open_monitoring_jmx") is True or item.get("open_monitoring_node") is True:
                  flags.append("open_monitoring_enabled_review")
              if item.get("customer_action_status"):
                  flags.append("customer_action_review")

              item["review_flags"] = sorted(set(flags))
              return item

          def lambda_handler(event, context):
              event = event if isinstance(event, dict) else {}
              excluded = csv_set(event.get("excluded_regions_csv", os.environ.get("EXCLUDED_REGIONS_CSV", "ap-northeast-3")))
              include_details = as_bool(event.get("include_details", os.environ.get("INCLUDE_CLUSTER_DETAILS", "true")), True)
              include_tags = as_bool(event.get("include_tags", os.environ.get("INCLUDE_TAGS", "true")), True)
              include_nodes = as_bool(event.get("include_nodes", os.environ.get("INCLUDE_NODES", "true")), True)
              max_clusters = as_int(event.get("max_clusters_per_region", os.environ.get("MAX_CLUSTERS_PER_REGION", "100")), 100)
              max_nodes = as_int(event.get("max_nodes_per_cluster", os.environ.get("MAX_NODES_PER_CLUSTER", "200")), 200)

              result = {
                  "module": "amazon_msk_cluster_cost_review",
                  "review_only": True,
                  "started_at": now(),
                  "ended_at": None,
                  "excluded_regions": sorted(excluded),
                  "totals": {
                      "regions_scanned": 0,
                      "regions_skipped": 0,
                      "clusters": 0,
                      "provisioned_clusters": 0,
                      "serverless_clusters": 0,
                      "active_clusters": 0,
                      "non_active_clusters": 0,
                      "total_broker_nodes": 0,
                      "total_ebs_storage_gib": 0,
                      "logging_enabled": 0,
                      "enhanced_monitoring": 0,
                      "open_monitoring_enabled": 0,
                      "regions_with_clusters": 0,
                      "regions_with_review_flags": 0,
                      "truncated_regions": 0,
                      "access_denied": 0,
                      "failed": 0
                  },
                  "regions": [],
                  "clusters": []
              }

              try:
                  regions = enabled_regions()
              except Exception as e:
                  result["fatal_error"] = str(e)
                  result["traceback"] = traceback.format_exc()
                  result["ended_at"] = now()
                  return result

              for region in regions:
                  if region in excluded:
                      result["totals"]["regions_skipped"] += 1
                      result["regions"].append({"region": region, "skipped": True, "reason": "excluded"})
                      continue

                  rec = {
                      "region": region,
                      "skipped": False,
                      "clusters": 0,
                      "provisioned_clusters": 0,
                      "serverless_clusters": 0,
                      "active_clusters": 0,
                      "non_active_clusters": 0,
                      "broker_nodes": 0,
                      "ebs_storage_gib": 0,
                      "logging_enabled": 0,
                      "enhanced_monitoring": 0,
                      "open_monitoring_enabled": 0,
                      "review_flags": [],
                      "truncated": False,
                      "error": None
                  }

                  try:
                      kafka = boto3.client("kafka", region_name=region)
                      clusters, truncated = get_clusters(kafka, max_clusters)
                      rec["truncated"] = truncated
                      if truncated:
                          result["totals"]["truncated_regions"] += 1

                      details = []
                      for c in clusters:
                          d = summarize_cluster(region, kafka, c, include_tags, include_nodes, max_nodes)
                          details.append(d)
                          rec["clusters"] += 1
                          result["totals"]["clusters"] += 1

                          if d.get("cluster_type") == "PROVISIONED":
                              rec["provisioned_clusters"] += 1
                              result["totals"]["provisioned_clusters"] += 1
                          if d.get("cluster_type") == "SERVERLESS":
                              rec["serverless_clusters"] += 1
                              result["totals"]["serverless_clusters"] += 1

                          if str(d.get("state") or "").upper() == "ACTIVE":
                              rec["active_clusters"] += 1
                              result["totals"]["active_clusters"] += 1
                          else:
                              rec["non_active_clusters"] += 1
                              result["totals"]["non_active_clusters"] += 1

                          nodes = d.get("number_of_broker_nodes") or 0
                          rec["broker_nodes"] += nodes
                          result["totals"]["total_broker_nodes"] += nodes

                          storage = d.get("ebs_volume_size_gib") or 0
                          if nodes and storage:
                              rec["ebs_storage_gib"] += nodes * storage
                              result["totals"]["total_ebs_storage_gib"] += nodes * storage

                          if d.get("logging_enabled") is True:
                              rec["logging_enabled"] += 1
                              result["totals"]["logging_enabled"] += 1
                          if d.get("enhanced_monitoring") and d.get("enhanced_monitoring") != "DEFAULT":
                              rec["enhanced_monitoring"] += 1
                              result["totals"]["enhanced_monitoring"] += 1
                          if d.get("open_monitoring_jmx") is True or d.get("open_monitoring_node") is True:
                              rec["open_monitoring_enabled"] += 1
                              result["totals"]["open_monitoring_enabled"] += 1

                          rec["review_flags"].extend(d.get("review_flags", []))

                      rec["review_flags"] = sorted(set(rec["review_flags"]))

                      if rec["clusters"] > 0:
                          result["totals"]["regions_with_clusters"] += 1
                      if rec["review_flags"]:
                          result["totals"]["regions_with_review_flags"] += 1

                      if include_details:
                          result["clusters"].extend(details)

                      result["totals"]["regions_scanned"] += 1
                      result["regions"].append(rec)

                  except ClientError as e:
                      code = e.response.get("Error", {}).get("Code", "ClientError")
                      rec["error"] = code
                      if code in ("AccessDenied", "AccessDeniedException", "UnauthorizedOperation"):
                          result["totals"]["access_denied"] += 1
                      else:
                          result["totals"]["failed"] += 1
                      result["regions"].append(rec)

                  except EndpointConnectionError:
                      rec["error"] = "EndpointConnectionError"
                      result["totals"]["failed"] += 1
                      result["regions"].append(rec)

                  except Exception as e:
                      rec["error"] = str(e)
                      rec["traceback"] = traceback.format_exc()
                      result["totals"]["failed"] += 1
                      result["regions"].append(rec)

              result["ended_at"] = now()
              return result

  MskCostReviewSchedule:
    Type: AWS::Events::Rule
    Condition: CreateSchedule
    Properties:
      Description: Scheduled review-only Amazon MSK cluster cost review.
      ScheduleExpression: !Ref ScheduleExpression
      State: ENABLED
      Targets:
        - Arn: !GetAtt MskCostReviewFunction.Arn
          Id: MskCostReviewFunctionTarget
          Input: '{"scheduled":true}'

  MskCostReviewSchedulePermission:
    Type: AWS::Lambda::Permission
    Condition: CreateSchedule
    Properties:
      FunctionName: !Ref MskCostReviewFunction
      Action: lambda:InvokeFunction
      Principal: events.amazonaws.com
      SourceArn: !GetAtt MskCostReviewSchedule.Arn

Outputs:
  ReviewFunctionName:
    Description: Lambda function name.
    Value: !Ref MskCostReviewFunction
  ReviewFunctionArn:
    Description: Lambda function ARN.
    Value: !GetAtt MskCostReviewFunction.Arn
  ReviewRoleArn:
    Description: IAM role ARN used by the review Lambda.
    Value: !GetAtt MskCostReviewRole.Arn
