AWSTemplateFormatVersion: '2010-09-09'
Description: Amazon MQ Broker 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
  IncludeBrokerDetails:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
  IncludeTags:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
  MaxBrokersPerRegion:
    Type: Number
    Default: 100
    MinValue: 1
    MaxValue: 200

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

Resources:
  AmazonMqCostReviewRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: [lambda.amazonaws.com]
            Action: [sts:AssumeRole]
      Policies:
        - PolicyName: amazon-mq-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: AmazonMqReadOnly
                Effect: Allow
                Action: [mq:ListBrokers, mq:DescribeBroker, mq:ListTags]
                Resource: '*'

  AmazonMqCostReviewFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Sub amazon-mq-cost-review-${AWS::StackName}
      Runtime: python3.12
      Handler: index.lambda_handler
      Role: !GetAtt AmazonMqCostReviewRole.Arn
      Timeout: 900
      MemorySize: 256
      Environment:
        Variables:
          EXCLUDED_REGIONS_CSV: !Ref ExcludedRegionsCsv
          INCLUDE_BROKER_DETAILS: !Ref IncludeBrokerDetails
          INCLUDE_TAGS: !Ref IncludeTags
          MAX_BROKERS_PER_REGION: !Ref MaxBrokersPerRegion
      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 list_brokers(mq, max_brokers):
              brokers = []
              token = None
              while True:
                  kwargs = {"MaxResults": min(max_brokers, 100)}
                  if token:
                      kwargs["NextToken"] = token
                  response = mq.list_brokers(**kwargs)
                  brokers.extend(response.get("BrokerSummaries", []))
                  if len(brokers) >= max_brokers:
                      return brokers[:max_brokers], True
                  token = response.get("NextToken")
                  if not token:
                      return brokers, False

          def log_summary(logs):
              logs = logs if isinstance(logs, dict) else {}
              return {
                  "audit": logs.get("Audit"),
                  "general": logs.get("General"),
                  "pending": logs.get("Pending")
              }

          def describe_one(region, mq, summary, include_tags):
              broker_id = summary.get("BrokerId")
              item = {
                  "region": region,
                  "broker_id": broker_id,
                  "broker_name": summary.get("BrokerName"),
                  "broker_state": summary.get("BrokerState"),
                  "deployment_mode": summary.get("DeploymentMode"),
                  "engine_type": summary.get("EngineType"),
                  "host_instance_type": summary.get("HostInstanceType"),
                  "created": str(summary.get("Created")) if summary.get("Created") else None,
                  "review_flags": []
              }

              try:
                  detail = mq.describe_broker(BrokerId=broker_id)
                  broker_arn = detail.get("BrokerArn")
                  item.update({
                      "broker_arn": broker_arn,
                      "engine_version": detail.get("EngineVersion"),
                      "auto_minor_version_upgrade": detail.get("AutoMinorVersionUpgrade"),
                      "deployment_mode": detail.get("DeploymentMode"),
                      "host_instance_type": detail.get("HostInstanceType"),
                      "publicly_accessible": detail.get("PubliclyAccessible"),
                      "storage_type": detail.get("StorageType"),
                      "subnet_ids_count": len(detail.get("SubnetIds", []) or []),
                      "security_groups_count": len(detail.get("SecurityGroups", []) or []),
                      "users_count": len(detail.get("Users", []) or []),
                      "logs": log_summary(detail.get("Logs")),
                      "maintenance_window_start_time": detail.get("MaintenanceWindowStartTime"),
                      "pending_engine_version": detail.get("PendingEngineVersion"),
                      "pending_host_instance_type": detail.get("PendingHostInstanceType"),
                      "actions_required_count": len(detail.get("ActionsRequired", []) or []),
                      "authentication_strategy": detail.get("AuthenticationStrategy")
                  })

                  if include_tags and broker_arn:
                      try:
                          tags = mq.list_tags(ResourceArn=broker_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")

              except ClientError as e:
                  item["describe_error"] = e.response.get("Error", {}).get("Code", "ClientError")
              except Exception as e:
                  item["describe_error"] = str(e)

              state = str(item.get("broker_state") or "").upper()
              deployment = str(item.get("deployment_mode") or "").upper()
              instance_type = str(item.get("host_instance_type") or "")
              logs = item.get("logs") or {}

              flags = []
              if state == "RUNNING":
                  flags.append("running_broker_review")
              elif state:
                  flags.append("broker_state_" + state.lower())

              if deployment in ("ACTIVE_STANDBY_MULTI_AZ", "CLUSTER_MULTI_AZ"):
                  flags.append("multi_az_broker_review")

              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("publicly_accessible") is True:
                  flags.append("publicly_accessible_review")

              if logs.get("audit") is True or logs.get("general") is True:
                  flags.append("broker_logging_enabled_review")

              if item.get("pending_engine_version") or item.get("pending_host_instance_type"):
                  flags.append("pending_change_review")

              if item.get("actions_required_count", 0) > 0:
                  flags.append("actions_required_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_BROKER_DETAILS", "true")), True)
              include_tags = as_bool(event.get("include_tags", os.environ.get("INCLUDE_TAGS", "true")), True)
              max_brokers = as_int(event.get("max_brokers_per_region", os.environ.get("MAX_BROKERS_PER_REGION", "100")), 100)

              result = {
                  "module": "amazon_mq_broker_cost_review",
                  "review_only": True,
                  "started_at": now(),
                  "ended_at": None,
                  "excluded_regions": sorted(excluded),
                  "totals": {
                      "regions_scanned": 0,
                      "regions_skipped": 0,
                      "brokers": 0,
                      "running_brokers": 0,
                      "stopped_brokers": 0,
                      "single_instance": 0,
                      "active_standby_multi_az": 0,
                      "cluster_multi_az": 0,
                      "publicly_accessible": 0,
                      "logging_enabled": 0,
                      "pending_changes": 0,
                      "actions_required": 0,
                      "regions_with_brokers": 0,
                      "regions_with_review_flags": 0,
                      "truncated_regions": 0,
                      "access_denied": 0,
                      "failed": 0
                  },
                  "regions": [],
                  "brokers": []
              }

              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,
                      "brokers": 0,
                      "running_brokers": 0,
                      "stopped_brokers": 0,
                      "single_instance": 0,
                      "active_standby_multi_az": 0,
                      "cluster_multi_az": 0,
                      "publicly_accessible": 0,
                      "logging_enabled": 0,
                      "pending_changes": 0,
                      "actions_required": 0,
                      "review_flags": [],
                      "truncated": False,
                      "error": None
                  }

                  try:
                      mq = boto3.client("mq", region_name=region)
                      summaries, truncated = list_brokers(mq, max_brokers)
                      rec["truncated"] = truncated
                      if truncated:
                          result["totals"]["truncated_regions"] += 1

                      details = []
                      for s in summaries:
                          d = describe_one(region, mq, s, include_tags)
                          details.append(d)
                          rec["brokers"] += 1
                          result["totals"]["brokers"] += 1

                          state = str(d.get("broker_state") or "").upper()
                          dep = str(d.get("deployment_mode") or "").upper()
                          logs = d.get("logs") or {}

                          if state == "RUNNING":
                              rec["running_brokers"] += 1
                              result["totals"]["running_brokers"] += 1
                          if state == "STOPPED":
                              rec["stopped_brokers"] += 1
                              result["totals"]["stopped_brokers"] += 1

                          if dep == "SINGLE_INSTANCE":
                              rec["single_instance"] += 1
                              result["totals"]["single_instance"] += 1
                          if dep == "ACTIVE_STANDBY_MULTI_AZ":
                              rec["active_standby_multi_az"] += 1
                              result["totals"]["active_standby_multi_az"] += 1
                          if dep == "CLUSTER_MULTI_AZ":
                              rec["cluster_multi_az"] += 1
                              result["totals"]["cluster_multi_az"] += 1

                          if d.get("publicly_accessible") is True:
                              rec["publicly_accessible"] += 1
                              result["totals"]["publicly_accessible"] += 1
                          if logs.get("audit") is True or logs.get("general") is True:
                              rec["logging_enabled"] += 1
                              result["totals"]["logging_enabled"] += 1
                          if d.get("pending_engine_version") or d.get("pending_host_instance_type"):
                              rec["pending_changes"] += 1
                              result["totals"]["pending_changes"] += 1
                          if d.get("actions_required_count", 0) > 0:
                              rec["actions_required"] += 1
                              result["totals"]["actions_required"] += 1

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

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

                      if rec["brokers"] > 0:
                          result["totals"]["regions_with_brokers"] += 1
                      if rec["review_flags"]:
                          result["totals"]["regions_with_review_flags"] += 1

                      if include_details:
                          result["brokers"].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

  AmazonMqCostReviewSchedule:
    Type: AWS::Events::Rule
    Condition: CreateSchedule
    Properties:
      Description: Scheduled review-only Amazon MQ broker cost review.
      ScheduleExpression: !Ref ScheduleExpression
      State: ENABLED
      Targets:
        - Arn: !GetAtt AmazonMqCostReviewFunction.Arn
          Id: AmazonMqCostReviewFunctionTarget
          Input: '{"scheduled":true}'

  AmazonMqCostReviewSchedulePermission:
    Type: AWS::Lambda::Permission
    Condition: CreateSchedule
    Properties:
      FunctionName: !Ref AmazonMqCostReviewFunction
      Action: lambda:InvokeFunction
      Principal: events.amazonaws.com
      SourceArn: !GetAtt AmazonMqCostReviewSchedule.Arn

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