AWSTemplateFormatVersion: '2010-09-09'
Description: >
  Amazon MQ / MSK / Kafka Connect Cost Review Remediator.
  Review-only inventory of Amazon MQ brokers, MSK provisioned/serverless clusters,
  and MSK Connect connectors across enabled Regions.

Parameters:
  ProjectName:
    Type: String
    Default: mq-msk-cost-review
  IncludeDetails:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
  ExcludedRegionsCsv:
    Type: String
    Default: ap-northeast-3
  ScheduleEnabled:
    Type: String
    Default: 'false'
    AllowedValues: ['true', 'false']
  ScheduleExpression:
    Type: String
    Default: rate(30 days)

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

Resources:
  ReviewRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal: { Service: lambda.amazonaws.com }
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: MqMskCostReviewReadOnlyPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action: ec2:DescribeRegions
                Resource: '*'
              - Effect: Allow
                Action:
                  - mq:ListBrokers
                  - mq:DescribeBroker
                  - kafka:ListClustersV2
                  - kafka:DescribeClusterV2
                  - kafkaconnect:ListConnectors
                  - kafkaconnect:DescribeConnector
                  - kafkaconnect:ListCustomPlugins
                  - kafkaconnect:ListWorkerConfigurations
                Resource: '*'

  ReviewFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Ref ProjectName
      Runtime: python3.12
      Handler: index.handler
      Timeout: 900
      MemorySize: 512
      Role: !GetAtt ReviewRole.Arn
      Environment:
        Variables:
          INCLUDE_DETAILS: !Ref IncludeDetails
          EXCLUDED_REGIONS_CSV: !Ref ExcludedRegionsCsv
      Code:
        ZipFile: |
          import json, os
          from datetime import datetime, timezone
          import boto3
          from botocore.exceptions import ClientError, EndpointConnectionError

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

          def boolish(v):
              return str(v).strip().lower() in ("1", "true", "yes", "y")

          def csv(v):
              return [x.strip() for x in (v or "").split(",") if x.strip()]

          def err(service, op, e):
              if isinstance(e, ClientError):
                  return {
                      "service": service,
                      "operation": op,
                      "error_code": e.response.get("Error", {}).get("Code"),
                      "error_message": e.response.get("Error", {}).get("Message")
                  }
              return {"service": service, "operation": op, "error_code": type(e).__name__, "error_message": str(e)}

          def denied(e):
              s = str(e).lower()
              if isinstance(e, ClientError):
                  s = (e.response.get("Error", {}).get("Code", "") + " " + e.response.get("Error", {}).get("Message", "")).lower()
              return "accessdenied" in s or "not authorized" in s or "explicit deny" in s or "service control policy" in s

          def safe(rr, svc, op, fn, default):
              try:
                  return fn()
              except EndpointConnectionError as e:
                  rr["failed"] += 1
                  rr["failures"].append(err(svc, op, e))
                  return default
              except ClientError as e:
                  if denied(e):
                      rr["access_denied"] += 1
                  else:
                      rr["failed"] += 1
                  rr["failures"].append(err(svc, op, e))
                  return default
              except Exception as e:
                  rr["failed"] += 1
                  rr["failures"].append(err(svc, op, e))
                  return default

          def enabled_regions(excluded):
              ec2 = boto3.client("ec2", region_name="us-east-1")
              regions = []
              for r in ec2.describe_regions(AllRegions=True).get("Regions", []):
                  status = r.get("OptInStatus")
                  name = r.get("RegionName")
                  if status in (None, "opt-in-not-required", "opted-in") and name not in excluded:
                      regions.append(name)
              return sorted(regions)

          def list_mq_brokers(client):
              out, token = [], None
              while True:
                  kw = {"MaxResults": 100}
                  if token:
                      kw["NextToken"] = token
                  resp = client.list_brokers(**kw)
                  out.extend(resp.get("BrokerSummaries", []) or [])
                  token = resp.get("NextToken")
                  if not token:
                      return out

          def list_msk_clusters(client):
              out, token = [], None
              while True:
                  kw = {"MaxResults": 100}
                  if token:
                      kw["NextToken"] = token
                  resp = client.list_clusters_v2(**kw)
                  out.extend(resp.get("ClusterInfoList", []) or [])
                  token = resp.get("NextToken")
                  if not token:
                      return out

          def list_connectors(client):
              out, token = [], None
              while True:
                  kw = {"maxResults": 100}
                  if token:
                      kw["nextToken"] = token
                  resp = client.list_connectors(**kw)
                  out.extend(resp.get("connectors", []) or [])
                  token = resp.get("nextToken")
                  if not token:
                      return out

          def list_custom_plugins(client):
              out, token = [], None
              while True:
                  kw = {"maxResults": 100}
                  if token:
                      kw["nextToken"] = token
                  resp = client.list_custom_plugins(**kw)
                  out.extend(resp.get("customPlugins", []) or [])
                  token = resp.get("nextToken")
                  if not token:
                      return out

          def list_worker_configs(client):
              out, token = [], None
              while True:
                  kw = {"maxResults": 100}
                  if token:
                      kw["nextToken"] = token
                  resp = client.list_worker_configurations(**kw)
                  out.extend(resp.get("workerConfigurations", []) or [])
                  token = resp.get("nextToken")
                  if not token:
                      return out

          def scan(region, details):
              rr = {
                  "region": region,
                  "mq_brokers": 0,
                  "mq_running_brokers": 0,
                  "mq_creation_in_progress_brokers": 0,
                  "mq_reboot_in_progress_brokers": 0,
                  "mq_publicly_accessible_brokers": 0,
                  "mq_single_instance_brokers": 0,
                  "mq_active_standby_brokers": 0,
                  "mq_cluster_multi_az_brokers": 0,
                  "mq_ebs_storage_brokers": 0,
                  "mq_efs_storage_brokers": 0,
                  "msk_clusters": 0,
                  "msk_provisioned_clusters": 0,
                  "msk_serverless_clusters": 0,
                  "msk_active_clusters": 0,
                  "msk_creating_clusters": 0,
                  "msk_broker_count": 0,
                  "msk_public_access_enabled_clusters": 0,
                  "msk_connectors": 0,
                  "msk_running_connectors": 0,
                  "msk_failed_connectors": 0,
                  "msk_connector_plugins": 0,
                  "msk_worker_configurations": 0,
                  "review_flags": [],
                  "mq_inventory": [],
                  "msk_inventory": [],
                  "connector_inventory": [],
                  "access_denied": 0,
                  "failed": 0,
                  "failures": []
              }

              mq = boto3.client("mq", region_name=region)
              brokers = safe(rr, "AmazonMQ", "ListBrokers", lambda: list_mq_brokers(mq), [])
              rr["mq_brokers"] = len(brokers)
              for b in brokers:
                  broker_id = b.get("BrokerId")
                  desc = safe(rr, "AmazonMQ", "DescribeBroker", lambda x=broker_id: mq.describe_broker(BrokerId=x), {})
                  state = desc.get("BrokerState") or b.get("BrokerState")
                  deploy = desc.get("DeploymentMode")
                  storage = desc.get("StorageType")
                  if state == "RUNNING":
                      rr["mq_running_brokers"] += 1
                  elif state == "CREATION_IN_PROGRESS":
                      rr["mq_creation_in_progress_brokers"] += 1
                  elif state == "REBOOT_IN_PROGRESS":
                      rr["mq_reboot_in_progress_brokers"] += 1
                  if desc.get("PubliclyAccessible"):
                      rr["mq_publicly_accessible_brokers"] += 1
                  if deploy == "SINGLE_INSTANCE":
                      rr["mq_single_instance_brokers"] += 1
                  elif deploy == "ACTIVE_STANDBY_MULTI_AZ":
                      rr["mq_active_standby_brokers"] += 1
                  elif deploy == "CLUSTER_MULTI_AZ":
                      rr["mq_cluster_multi_az_brokers"] += 1
                  if storage == "EBS":
                      rr["mq_ebs_storage_brokers"] += 1
                  elif storage == "EFS":
                      rr["mq_efs_storage_brokers"] += 1
                  if details:
                      rr["mq_inventory"].append({
                          "broker_id": broker_id,
                          "broker_name": desc.get("BrokerName") or b.get("BrokerName"),
                          "broker_state": state,
                          "engine_type": desc.get("EngineType"),
                          "engine_version": desc.get("EngineVersion"),
                          "host_instance_type": desc.get("HostInstanceType"),
                          "deployment_mode": deploy,
                          "storage_type": storage,
                          "publicly_accessible": desc.get("PubliclyAccessible"),
                          "created": desc.get("Created")
                      })

              kafka = boto3.client("kafka", region_name=region)
              clusters = safe(rr, "MSK", "ListClustersV2", lambda: list_msk_clusters(kafka), [])
              rr["msk_clusters"] = len(clusters)
              for c in clusters:
                  arn = c.get("ClusterArn")
                  ctype = c.get("ClusterType")
                  state = c.get("State")
                  detail = safe(rr, "MSK", "DescribeClusterV2", lambda a=arn: kafka.describe_cluster_v2(ClusterArn=a).get("ClusterInfo", {}), {})
                  if ctype == "PROVISIONED":
                      rr["msk_provisioned_clusters"] += 1
                  elif ctype == "SERVERLESS":
                      rr["msk_serverless_clusters"] += 1
                  if state == "ACTIVE":
                      rr["msk_active_clusters"] += 1
                  elif state == "CREATING":
                      rr["msk_creating_clusters"] += 1
                  provisioned = detail.get("Provisioned") or {}
                  broker_nodes = int((provisioned.get("NumberOfBrokerNodes") or 0))
                  rr["msk_broker_count"] += broker_nodes
                  public_access = (((provisioned.get("BrokerNodeGroupInfo") or {}).get("ConnectivityInfo") or {}).get("PublicAccess") or {}).get("Type")
                  if public_access and public_access != "DISABLED":
                      rr["msk_public_access_enabled_clusters"] += 1
                  if details:
                      rr["msk_inventory"].append({
                          "cluster_name": c.get("ClusterName"),
                          "cluster_arn": arn,
                          "cluster_type": ctype,
                          "state": state,
                          "creation_time": c.get("CreationTime"),
                          "number_of_broker_nodes": broker_nodes,
                          "broker_instance_type": (provisioned.get("BrokerNodeGroupInfo") or {}).get("InstanceType"),
                          "kafka_version": provisioned.get("CurrentBrokerSoftwareInfo", {}).get("KafkaVersion"),
                          "public_access": public_access
                      })

              kc = boto3.client("kafkaconnect", region_name=region)
              connectors = safe(rr, "MSKConnect", "ListConnectors", lambda: list_connectors(kc), [])
              rr["msk_connectors"] = len(connectors)
              for con in connectors:
                  state = con.get("connectorState")
                  if state == "RUNNING":
                      rr["msk_running_connectors"] += 1
                  elif state == "FAILED":
                      rr["msk_failed_connectors"] += 1
                  if details:
                      arn = con.get("connectorArn")
                      desc = safe(rr, "MSKConnect", "DescribeConnector", lambda a=arn: kc.describe_connector(connectorArn=a), {})
                      rr["connector_inventory"].append({
                          "connector_name": con.get("connectorName"),
                          "connector_arn": arn,
                          "connector_state": state,
                          "kafka_cluster_type": "apache_kafka_cluster" if desc.get("kafkaCluster", {}).get("apacheKafkaCluster") else "msk_cluster" if desc.get("kafkaCluster", {}).get("apacheKafkaCluster") is None else "unknown",
                          "capacity": desc.get("capacity"),
                          "kafka_connect_version": desc.get("kafkaConnectVersion"),
                          "creation_time": con.get("creationTime")
                      })

              plugins = safe(rr, "MSKConnect", "ListCustomPlugins", lambda: list_custom_plugins(kc), [])
              rr["msk_connector_plugins"] = len(plugins)
              workers = safe(rr, "MSKConnect", "ListWorkerConfigurations", lambda: list_worker_configs(kc), [])
              rr["msk_worker_configurations"] = len(workers)

              if rr["mq_running_brokers"] > 0:
                  rr["review_flags"].append("running_amazon_mq_brokers_present")
              if rr["mq_publicly_accessible_brokers"] > 0:
                  rr["review_flags"].append("publicly_accessible_amazon_mq_brokers_present")
              if rr["msk_active_clusters"] > 0:
                  rr["review_flags"].append("active_msk_clusters_present")
              if rr["msk_broker_count"] > 0:
                  rr["review_flags"].append("msk_broker_nodes_present")
              if rr["msk_public_access_enabled_clusters"] > 0:
                  rr["review_flags"].append("msk_public_access_enabled")
              if rr["msk_running_connectors"] > 0:
                  rr["review_flags"].append("running_msk_connectors_present")
              if rr["msk_failed_connectors"] > 0:
                  rr["review_flags"].append("failed_msk_connectors_present")

              return rr

          def handler(event, context):
              event = event or {}
              excluded = set(csv(event.get("excluded_regions_csv") or os.environ.get("EXCLUDED_REGIONS_CSV", "ap-northeast-3")))
              details = boolish(event.get("include_details", os.environ.get("INCLUDE_DETAILS", "true")))

              totals = {k:0 for k in [
                  "regions_scanned", "mq_brokers", "mq_running_brokers", "mq_creation_in_progress_brokers",
                  "mq_reboot_in_progress_brokers", "mq_publicly_accessible_brokers", "mq_single_instance_brokers",
                  "mq_active_standby_brokers", "mq_cluster_multi_az_brokers", "mq_ebs_storage_brokers",
                  "mq_efs_storage_brokers", "msk_clusters", "msk_provisioned_clusters", "msk_serverless_clusters",
                  "msk_active_clusters", "msk_creating_clusters", "msk_broker_count", "msk_public_access_enabled_clusters",
                  "msk_connectors", "msk_running_connectors", "msk_failed_connectors", "msk_connector_plugins",
                  "msk_worker_configurations", "regions_with_mq_brokers", "regions_with_msk_clusters",
                  "regions_with_msk_connectors", "regions_with_review_flags", "access_denied", "failed"
              ]}

              output = {
                  "started_at": now(),
                  "review_only": True,
                  "include_details": details,
                  "excluded_regions": sorted(excluded),
                  "totals": totals,
                  "region_results": [],
                  "failures": []
              }

              for region in enabled_regions(excluded):
                  rr = scan(region, details)
                  for k in totals:
                      if k.startswith("regions_with_") or k == "regions_scanned":
                          continue
                      totals[k] += int(rr.get(k, 0) or 0)
                  totals["regions_scanned"] += 1
                  if rr["mq_brokers"] > 0:
                      totals["regions_with_mq_brokers"] += 1
                  if rr["msk_clusters"] > 0:
                      totals["regions_with_msk_clusters"] += 1
                  if rr["msk_connectors"] > 0:
                      totals["regions_with_msk_connectors"] += 1
                  if rr["review_flags"]:
                      totals["regions_with_review_flags"] += 1
                  output["failures"].extend([dict(f, region=region) for f in rr.get("failures", [])])
                  output["region_results"].append(rr)

              output["ended_at"] = now()
              print(json.dumps(output, indent=2, default=str))
              return json.loads(json.dumps(output, default=str))

  ReviewScheduleRule:
    Type: AWS::Events::Rule
    Condition: IsScheduleEnabled
    Properties:
      Name: !Sub '${ProjectName}-schedule'
      ScheduleExpression: !Ref ScheduleExpression
      State: ENABLED
      Targets:
        - Arn: !GetAtt ReviewFunction.Arn
          Id: MqMskCostReviewTarget
          Input: '{"source":"eventbridge-schedule"}'

  ReviewSchedulePermission:
    Type: AWS::Lambda::Permission
    Condition: IsScheduleEnabled
    Properties:
      FunctionName: !Ref ReviewFunction
      Action: lambda:InvokeFunction
      Principal: events.amazonaws.com
      SourceArn: !GetAtt ReviewScheduleRule.Arn

Outputs:
  ReviewFunctionName:
    Description: Lambda function name.
    Value: !Ref ReviewFunction
  RemediatorFunctionName:
    Description: Lambda function name alias.
    Value: !Ref ReviewFunction
  ReviewFunctionArn:
    Description: Lambda function ARN.
    Value: !GetAtt ReviewFunction.Arn
  ReviewOnly:
    Description: This module is intentionally review-only.
    Value: 'true'
