AWSTemplateFormatVersion: '2010-09-09'
Description: >
  SageMaker Notebook and Endpoint Review Remediator.
  Purpose-specific module for reviewing SageMaker notebook instances, endpoints,
  models, endpoint configs, training jobs, processing jobs, and transform jobs.
  Optional remediation can stop explicitly approved old notebook instances only.
  It does not delete SageMaker resources.

Parameters:
  ProjectName:
    Type: String
    Default: sagemaker-review
    Description: Short base name used for created resources.

  TargetRegionsCsv:
    Type: String
    Default: all-enabled
    Description: Comma-separated AWS Regions to scan, or all-enabled to scan all enabled Regions.

  ExcludedRegionsCsv:
    Type: String
    Default: ap-northeast-3
    Description: Comma-separated AWS Regions to exclude from scanning/remediation.

  DryRun:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
    Description: When true, only report what would happen. When false, perform eligible approved notebook stop actions.

  ActionMode:
    Type: String
    Default: review-only
    AllowedValues:
      - review-only
      - stop-approved-notebooks
    Description: review-only never changes resources. stop-approved-notebooks stops eligible approved notebook instances only.

  MinimumNotebookAgeDays:
    Type: Number
    Default: 30
    Description: Notebook instance must be at least this many days old before stop consideration.

  RequireApprovalTag:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
    Description: Notebook stop requires ApprovalTagKey=ApprovalTagValue when true.

  ApprovalTagKey:
    Type: String
    Default: CostOptApproved
    Description: Tag key required when RequireApprovalTag is true.

  ApprovalTagValue:
    Type: String
    Default: 'true'
    Description: Tag value required when RequireApprovalTag is true.

  IgnoreTagKeysCsv:
    Type: String
    Default: CostOptIgnore,DoNotStop,DoNotModify,ComplianceHold,LegalHold,Production,prod
    Description: Comma-separated tag keys that prevent notebook stop.

  IgnoreTagPairsCsv:
    Type: String
    Default: CostOptIgnore=true,DoNotStop=true,DoNotModify=true,ComplianceHold=true,LegalHold=true,Environment=Production,Environment=Prod
    Description: Comma-separated key=value tag pairs that prevent notebook stop.

  IncludeNotebookNamePrefixesCsv:
    Type: String
    Default: ''
    Description: Optional comma-separated notebook name prefixes to include. Leave blank to include all.

  ExcludeNotebookNamePrefixesCsv:
    Type: String
    Default: ''
    Description: Optional comma-separated notebook name prefixes to exclude.

  MaxNotebooksPerRun:
    Type: Number
    Default: 25
    Description: Safety limit. 0 means no limit. Recommended actual stop batch is 1 to 5.

  ContinueOnAccessDenied:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
    Description: Continue scanning other Regions/resources when blocked by IAM/SCP.

  ScheduleEnabled:
    Type: String
    Default: 'false'
    AllowedValues: ['true', 'false']
    Description: Whether to create an enabled EventBridge schedule.

  ScheduleExpression:
    Type: String
    Default: rate(30 days)
    Description: EventBridge schedule expression used when ScheduleEnabled is true.

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

Resources:
  RemediatorRole:
    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: SageMakerReviewPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Sid: DiscoverRegions
                Effect: Allow
                Action:
                  - ec2:DescribeRegions
                Resource: '*'
              - Sid: ReadSageMaker
                Effect: Allow
                Action:
                  - sagemaker:ListNotebookInstances
                  - sagemaker:DescribeNotebookInstance
                  - sagemaker:ListEndpoints
                  - sagemaker:DescribeEndpoint
                  - sagemaker:ListModels
                  - sagemaker:ListEndpointConfigs
                  - sagemaker:ListTrainingJobs
                  - sagemaker:ListProcessingJobs
                  - sagemaker:ListTransformJobs
                  - sagemaker:ListTags
                Resource: '*'
              - Sid: StopApprovedNotebookInstances
                Effect: Allow
                Action:
                  - sagemaker:StopNotebookInstance
                Resource: '*'

  RemediatorFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Ref ProjectName
      Description: SageMaker review Lambda with guarded optional stop for approved notebook instances only.
      Runtime: python3.12
      Handler: index.handler
      Timeout: 900
      MemorySize: 512
      Role: !GetAtt RemediatorRole.Arn
      Environment:
        Variables:
          TARGET_REGIONS_CSV: !Ref TargetRegionsCsv
          EXCLUDED_REGIONS_CSV: !Ref ExcludedRegionsCsv
          DRY_RUN: !Ref DryRun
          ACTION_MODE: !Ref ActionMode
          MINIMUM_NOTEBOOK_AGE_DAYS: !Ref MinimumNotebookAgeDays
          REQUIRE_APPROVAL_TAG: !Ref RequireApprovalTag
          APPROVAL_TAG_KEY: !Ref ApprovalTagKey
          APPROVAL_TAG_VALUE: !Ref ApprovalTagValue
          IGNORE_TAG_KEYS_CSV: !Ref IgnoreTagKeysCsv
          IGNORE_TAG_PAIRS_CSV: !Ref IgnoreTagPairsCsv
          INCLUDE_NOTEBOOK_NAME_PREFIXES_CSV: !Ref IncludeNotebookNamePrefixesCsv
          EXCLUDE_NOTEBOOK_NAME_PREFIXES_CSV: !Ref ExcludeNotebookNamePrefixesCsv
          MAX_NOTEBOOKS_PER_RUN: !Ref MaxNotebooksPerRun
          CONTINUE_ON_ACCESS_DENIED: !Ref ContinueOnAccessDenied
      Code:
        ZipFile: |
          import json
          import os
          import traceback
          from datetime import datetime, timezone

          import boto3
          from botocore.exceptions import ClientError

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

          def split_csv(value):
              return [x.strip() for x in str(value or '').split(',') if x.strip()]

          def env_bool(name, default=False):
              return os.environ.get(name, str(default)).strip().lower() in ('1', 'true', 'yes', 'y')

          def env_int(name, default=0):
              try:
                  return int(os.environ.get(name, str(default)).strip())
              except Exception:
                  return default

          def parse_tag_pairs(value):
              pairs = {}
              for item in split_csv(value):
                  if '=' not in item:
                      continue
                  key, val = item.split('=', 1)
                  if key.strip():
                      pairs[key.strip()] = val.strip()
              return pairs

          def age_days(dt):
              if not dt:
                  return None
              if dt.tzinfo is None:
                  dt = dt.replace(tzinfo=timezone.utc)
              return (now() - dt).days

          def enabled_regions(excluded):
              requested = os.environ.get('TARGET_REGIONS_CSV', 'all-enabled').strip()
              if requested and requested.lower() != 'all-enabled':
                  regions = split_csv(requested)
              else:
                  ec2 = boto3.client('ec2', region_name=os.environ.get('AWS_REGION', 'us-east-1'))
                  response = ec2.describe_regions(AllRegions=False)
                  regions = sorted([r['RegionName'] for r in response.get('Regions', [])])
              return [r for r in regions if r not in excluded]

          def paginate(client, operation_name, result_key, **kwargs):
              items = []
              try:
                  paginator = client.get_paginator(operation_name)
                  for page in paginator.paginate(**kwargs):
                      items.extend(page.get(result_key, []) or [])
                  return items
              except Exception:
                  request = dict(kwargs)
                  while True:
                      response = getattr(client, operation_name)(**request)
                      items.extend(response.get(result_key, []) or [])
                      token = response.get('NextToken')
                      if not token:
                          break
                      request['NextToken'] = token
                  return items

          def tags_to_dict(tag_list):
              return {t.get('Key'): t.get('Value', '') for t in tag_list or [] if t.get('Key')}

          def safe_tags(client, arn):
              if not arn:
                  return {}, None
              try:
                  response = client.list_tags(ResourceArn=arn)
                  return tags_to_dict(response.get('Tags', [])), None
              except ClientError as e:
                  return {}, e

          def has_ignore_tag(tags, ignore_keys, ignore_pairs):
              for key in ignore_keys:
                  if key in tags:
                      return True, 'ignore_tag_key=' + key
              for key, value in ignore_pairs.items():
                  if tags.get(key) == value:
                      return True, 'ignore_tag_pair=' + key + '=' + value
              return False, ''

          def has_approval_tag(tags, settings):
              if not settings['require_approval_tag']:
                  return True
              return tags.get(settings['approval_tag_key']) == settings['approval_tag_value']

          def name_included(value, prefixes):
              if not prefixes:
                  return True
              return any((value or '').startswith(prefix) for prefix in prefixes)

          def name_excluded(value, prefixes):
              if not prefixes:
                  return False
              return any((value or '').startswith(prefix) for prefix in prefixes)

          def count_list(client, op, key):
              try:
                  return len(paginate(client, op, key))
              except ClientError as e:
                  raise e

          def endpoint_inventory(endpoint):
              return {
                  'endpoint_name': endpoint.get('EndpointName'),
                  'endpoint_arn': endpoint.get('EndpointArn'),
                  'status': endpoint.get('EndpointStatus'),
                  'creation_time': str(endpoint.get('CreationTime')) if endpoint.get('CreationTime') else None,
                  'last_modified_time': str(endpoint.get('LastModifiedTime')) if endpoint.get('LastModifiedTime') else None
              }

          def scan_region(region, settings):
              client = boto3.client('sagemaker', region_name=region)
              result = {
                  'region': region,
                  'notebook_instances': 0,
                  'in_service_notebooks': 0,
                  'stopped_notebooks': 0,
                  'endpoints': 0,
                  'in_service_endpoints': 0,
                  'models': 0,
                  'endpoint_configs': 0,
                  'training_jobs': 0,
                  'processing_jobs': 0,
                  'transform_jobs': 0,
                  'old_in_service_notebooks': 0,
                  'eligible_notebooks': 0,
                  'would_stop_notebooks': 0,
                  'stopped_notebooks_by_action': 0,
                  'skipped_not_running': 0,
                  'skipped_too_new': 0,
                  'skipped_include_prefix': 0,
                  'skipped_exclude_prefix': 0,
                  'skipped_missing_approval_tag': 0,
                  'skipped_ignore_tag': 0,
                  'tag_read_access_denied': 0,
                  'access_denied': 0,
                  'failed': 0,
                  'failures': [],
                  'notebook_inventory': [],
                  'endpoint_inventory': [],
                  'candidates': [],
                  'stopped_notebook_names': []
              }

              try:
                  notebooks = paginate(client, 'list_notebook_instances', 'NotebookInstances')
                  endpoints = paginate(client, 'list_endpoints', 'Endpoints')
                  result['models'] = count_list(client, 'list_models', 'Models')
                  result['endpoint_configs'] = count_list(client, 'list_endpoint_configs', 'EndpointConfigs')
                  result['training_jobs'] = count_list(client, 'list_training_jobs', 'TrainingJobSummaries')
                  result['processing_jobs'] = count_list(client, 'list_processing_jobs', 'ProcessingJobSummaries')
                  result['transform_jobs'] = count_list(client, 'list_transform_jobs', 'TransformJobSummaries')
              except ClientError as e:
                  code = e.response.get('Error', {}).get('Code', 'Unknown')
                  message = e.response.get('Error', {}).get('Message', str(e))
                  result['failures'].append({'operation': 'ListSageMakerResources', 'error_code': code, 'error_message': message})
                  if code in ('AccessDenied', 'AccessDeniedException', 'UnauthorizedOperation') or 'explicit deny' in message.lower():
                      result['access_denied'] += 1
                  else:
                      result['failed'] += 1
                  if not settings['continue_on_access_denied']:
                      raise
                  return result

              result['notebook_instances'] = len(notebooks)
              result['endpoints'] = len(endpoints)

              for endpoint in endpoints:
                  result['endpoint_inventory'].append(endpoint_inventory(endpoint))
                  if endpoint.get('EndpointStatus') == 'InService':
                      result['in_service_endpoints'] += 1

              for nb in notebooks:
                  if settings['max_notebooks_per_run'] and (result['would_stop_notebooks'] + result['stopped_notebooks_by_action']) >= settings['max_notebooks_per_run']:
                      result['stopped_due_to_max_notebooks_per_run'] = True
                      return result

                  name = nb.get('NotebookInstanceName')
                  status = nb.get('NotebookInstanceStatus')
                  created = nb.get('CreationTime')
                  modified = nb.get('LastModifiedTime')
                  arn = nb.get('NotebookInstanceArn')
                  nb_age = age_days(created)
                  tags, tag_error = safe_tags(client, arn)

                  item = {
                      'notebook_instance_name': name,
                      'notebook_instance_arn': arn,
                      'status': status,
                      'instance_type': nb.get('InstanceType'),
                      'creation_time': str(created) if created else None,
                      'last_modified_time': str(modified) if modified else None,
                      'age_days': nb_age,
                      'tags': tags
                  }

                  if tag_error:
                      code = tag_error.response.get('Error', {}).get('Code', 'Unknown')
                      message = tag_error.response.get('Error', {}).get('Message', str(tag_error))
                      item['tag_read_error'] = {'error_code': code, 'error_message': message}
                      if code in ('AccessDenied', 'AccessDeniedException') or 'explicit deny' in message.lower():
                          result['tag_read_access_denied'] += 1
                      else:
                          result['failed'] += 1

                  if status == 'InService':
                      result['in_service_notebooks'] += 1
                  elif status == 'Stopped':
                      result['stopped_notebooks'] += 1

                  result['notebook_inventory'].append(item)

                  if not name_included(name, settings['include_notebook_prefixes']):
                      result['skipped_include_prefix'] += 1
                      continue

                  if name_excluded(name, settings['exclude_notebook_prefixes']):
                      result['skipped_exclude_prefix'] += 1
                      item['action'] = 'skipped_exclude_prefix'
                      result['candidates'].append(item)
                      continue

                  if status != 'InService':
                      result['skipped_not_running'] += 1
                      item['action'] = 'skipped_not_running'
                      result['candidates'].append(item)
                      continue

                  if nb_age is None or nb_age < settings['minimum_notebook_age_days']:
                      result['skipped_too_new'] += 1
                      item['action'] = 'skipped_too_new'
                      result['candidates'].append(item)
                      continue

                  result['old_in_service_notebooks'] += 1

                  ignore, reason = has_ignore_tag(tags, settings['ignore_tag_keys'], settings['ignore_tag_pairs'])
                  if ignore:
                      result['skipped_ignore_tag'] += 1
                      item['action'] = 'skipped_ignore_tag'
                      item['reason'] = reason
                      result['candidates'].append(item)
                      continue

                  if not has_approval_tag(tags, settings):
                      result['skipped_missing_approval_tag'] += 1
                      item['action'] = 'skipped_missing_approval_tag'
                      item['required_tag'] = settings['approval_tag_key'] + '=' + settings['approval_tag_value']
                      result['candidates'].append(item)
                      continue

                  result['eligible_notebooks'] += 1

                  if settings['action_mode'] != 'stop-approved-notebooks':
                      item['action'] = 'review_only_candidate'
                      result['candidates'].append(item)
                      continue

                  if settings['dry_run']:
                      result['would_stop_notebooks'] += 1
                      item['action'] = 'would_stop_notebook'
                      result['candidates'].append(item)
                      continue

                  try:
                      client.stop_notebook_instance(NotebookInstanceName=name)
                      result['stopped_notebooks_by_action'] += 1
                      item['action'] = 'stop_notebook'
                      result['stopped_notebook_names'].append(name)
                  except ClientError as e:
                      code = e.response.get('Error', {}).get('Code', 'Unknown')
                      message = e.response.get('Error', {}).get('Message', str(e))
                      result['failures'].append({'notebook_instance_name': name, 'operation': 'StopNotebookInstance', 'error_code': code, 'error_message': message})
                      if code in ('AccessDenied', 'AccessDeniedException') or 'explicit deny' in message.lower():
                          result['access_denied'] += 1
                      else:
                          result['failed'] += 1
                      if not settings['continue_on_access_denied']:
                          raise

              return result

          def handler(event, context):
              settings = {
                  'excluded_regions': split_csv(os.environ.get('EXCLUDED_REGIONS_CSV', '')),
                  'dry_run': env_bool('DRY_RUN', True),
                  'action_mode': os.environ.get('ACTION_MODE', 'review-only'),
                  'minimum_notebook_age_days': env_int('MINIMUM_NOTEBOOK_AGE_DAYS', 30),
                  'require_approval_tag': env_bool('REQUIRE_APPROVAL_TAG', True),
                  'approval_tag_key': os.environ.get('APPROVAL_TAG_KEY', 'CostOptApproved'),
                  'approval_tag_value': os.environ.get('APPROVAL_TAG_VALUE', 'true'),
                  'ignore_tag_keys': split_csv(os.environ.get('IGNORE_TAG_KEYS_CSV', '')),
                  'ignore_tag_pairs': parse_tag_pairs(os.environ.get('IGNORE_TAG_PAIRS_CSV', '')),
                  'include_notebook_prefixes': split_csv(os.environ.get('INCLUDE_NOTEBOOK_NAME_PREFIXES_CSV', '')),
                  'exclude_notebook_prefixes': split_csv(os.environ.get('EXCLUDE_NOTEBOOK_NAME_PREFIXES_CSV', '')),
                  'max_notebooks_per_run': env_int('MAX_NOTEBOOKS_PER_RUN', 25),
                  'continue_on_access_denied': env_bool('CONTINUE_ON_ACCESS_DENIED', True)
              }

              if isinstance(event, dict):
                  if 'dry_run' in event:
                      settings['dry_run'] = bool(event['dry_run'])
                  if 'action_mode' in event:
                      settings['action_mode'] = str(event['action_mode'])
                  if 'require_approval_tag' in event:
                      settings['require_approval_tag'] = bool(event['require_approval_tag'])
                  if 'minimum_notebook_age_days' in event:
                      settings['minimum_notebook_age_days'] = int(event['minimum_notebook_age_days'])
                  if 'max_notebooks_per_run' in event:
                      settings['max_notebooks_per_run'] = int(event['max_notebooks_per_run'])
                  if 'target_regions' in event and event['target_regions']:
                      os.environ['TARGET_REGIONS_CSV'] = ','.join(event['target_regions'])

              if settings['action_mode'] not in ('review-only', 'stop-approved-notebooks'):
                  settings['action_mode'] = 'review-only'

              summary = {
                  'started_at': now().isoformat(),
                  'dry_run': settings['dry_run'],
                  'action_mode': settings['action_mode'],
                  'minimum_notebook_age_days': settings['minimum_notebook_age_days'],
                  'require_approval_tag': settings['require_approval_tag'],
                  'approval_tag': settings['approval_tag_key'] + '=' + settings['approval_tag_value'],
                  'max_notebooks_per_run': settings['max_notebooks_per_run'],
                  'excluded_regions': settings['excluded_regions'],
                  'regions': [],
                  'totals': {
                      'notebook_instances': 0,
                      'in_service_notebooks': 0,
                      'stopped_notebooks': 0,
                      'endpoints': 0,
                      'in_service_endpoints': 0,
                      'models': 0,
                      'endpoint_configs': 0,
                      'training_jobs': 0,
                      'processing_jobs': 0,
                      'transform_jobs': 0,
                      'old_in_service_notebooks': 0,
                      'eligible_notebooks': 0,
                      'would_stop_notebooks': 0,
                      'stopped_notebooks_by_action': 0,
                      'skipped_not_running': 0,
                      'skipped_too_new': 0,
                      'skipped_missing_approval_tag': 0,
                      'skipped_ignore_tag': 0,
                      'tag_read_access_denied': 0,
                      'access_denied': 0,
                      'failed': 0
                  },
                  'region_results': []
              }

              try:
                  regions = enabled_regions(settings['excluded_regions'])
                  summary['regions'] = regions
                  for region in regions:
                      try:
                          rr = scan_region(region, settings)
                      except Exception as e:
                          rr = {
                              'region': region,
                              'notebook_instances': 0,
                              'in_service_notebooks': 0,
                              'stopped_notebooks': 0,
                              'endpoints': 0,
                              'in_service_endpoints': 0,
                              'models': 0,
                              'endpoint_configs': 0,
                              'training_jobs': 0,
                              'processing_jobs': 0,
                              'transform_jobs': 0,
                              'old_in_service_notebooks': 0,
                              'eligible_notebooks': 0,
                              'would_stop_notebooks': 0,
                              'stopped_notebooks_by_action': 0,
                              'skipped_not_running': 0,
                              'skipped_too_new': 0,
                              'skipped_missing_approval_tag': 0,
                              'skipped_ignore_tag': 0,
                              'tag_read_access_denied': 0,
                              'access_denied': 0,
                              'failed': 1,
                              'failures': [{'error_code': type(e).__name__, 'error_message': str(e), 'trace': traceback.format_exc()}],
                              'notebook_inventory': [],
                              'endpoint_inventory': [],
                              'candidates': []
                          }
                      summary['region_results'].append(rr)
                      for key in summary['totals']:
                          summary['totals'][key] += int(rr.get(key, 0) or 0)
              finally:
                  summary['ended_at'] = now().isoformat()
                  print(json.dumps(summary, indent=2, default=str))
              return summary

  RemediatorScheduleRule:
    Type: AWS::Events::Rule
    Condition: IsScheduleEnabled
    Properties:
      Name: !Sub '${ProjectName}-schedule'
      Description: Scheduled run for SageMaker Review Remediator.
      ScheduleExpression: !Ref ScheduleExpression
      State: ENABLED
      Targets:
        - Arn: !GetAtt RemediatorFunction.Arn
          Id: SageMakerReviewTarget
          Input: '{"source":"eventbridge-schedule","dry_run":true,"action_mode":"review-only"}'

  RemediatorSchedulePermission:
    Type: AWS::Lambda::Permission
    Condition: IsScheduleEnabled
    Properties:
      FunctionName: !Ref RemediatorFunction
      Action: lambda:InvokeFunction
      Principal: events.amazonaws.com
      SourceArn: !GetAtt RemediatorScheduleRule.Arn

Outputs:
  RemediatorFunctionName:
    Description: Lambda function name.
    Value: !Ref RemediatorFunction

  RemediatorFunctionArn:
    Description: Lambda function ARN.
    Value: !GetAtt RemediatorFunction.Arn

  DryRunSetting:
    Description: Current DryRun setting.
    Value: !Ref DryRun

  ActionModeSetting:
    Description: Current ActionMode setting.
    Value: !Ref ActionMode
