AWSTemplateFormatVersion: "2010-09-09"
Transform: AWS::Serverless-2016-10-31
Parameters:
  ResourceUsageThreshold:
    Type: Number
    Default: 3
  SecurityGroupIds:
    Type: CommaDelimitedList
    Description: 'Security Groups for Lambda function'
  SubnetIds:
    Type: CommaDelimitedList
    Description: 'Subnets for Lambda function'
Resources:
  SageMakerNotebookInstanceMlt3mediumUsageAlarm:
    Type: 'AWS::CloudWatch::Alarm'
    Properties:
      ActionsEnabled: True
      AlarmActions:
        - !GetAtt UpdateQuotasTopic.TopicArn
      AlarmDescription: 'ml.t3.medium for SageMaker notebook instance usage above 50% alarm'
      ComparisonOperator: 'GreaterThanOrEqualToThreshold'
      EvaluationPeriods: 10
      Metrics:
        - Id: usage_data
          MetricStat:
            Metric:
              Dimensions:
              - Name: Class
                Value: None
              - Name: Resource
                Value: notebook-instance/ml.t3.medium
              - Name: Service
                Value: SageMaker
              - Name: Type
                Value: Resource
              MetricName: ResourceCount
              Namespace: AWS/Usage
            Period: 3600
            Stat: Maximum
          ReturnData: false
        - Expression: (usage_data/SERVICE_QUOTA(usage_data))*100
          Id: pct_utilization
          Label: '% Utilization'
          ReturnData: true
      Threshold: !Ref ResourceUsageThreshold
  DefaultKey:
    Type: 'AWS::KMS::Key'
    Properties:
      EnableKeyRotation: True
      KeyPolicy:
        Version: 2012-10-17
        Id: key-default-1
        Statement:
          - Sid: Enable IAM User Permissions
            Effect: Allow
            Principal:
              AWS: !Sub 'arn:aws:iam::${AWS::AccountId}:root'
            Action: 'kms:*'
            Resource: '*'
          - Sid: Allow administration of the key
            Effect: Allow
            Principal:
              AWS: !Sub 'arn:aws:iam::${AWS::AccountId}:root'
            Action:
              - 'kms:Create*'
              - 'kms:Describe*'
              - 'kms:Enable*'
              - 'kms:List*'
              - 'kms:Put*'
              - 'kms:Update*'
              - 'kms:Revoke*'
              - 'kms:Disable*'
              - 'kms:Get*'
              - 'kms:Delete*'
              - 'kms:ScheduleKeyDeletion'
              - 'kms:CancelKeyDeletion'
            Resource: '*'
          - Sid: Allow use of the key
            Effect: Allow
            Principal:
              AWS: !Sub 'arn:aws:iam::${AWS::AccountId}:root'
            Action:
              - 'kms:DescribeKey'
              - 'kms:Encrypt'
              - 'kms:Decrypt'
              - 'kms:ReEncrypt*'
              - 'kms:GenerateDataKey'
              - 'kms:GenerateDataKeyWithoutPlaintext'
            Resource: '*'
  UpdateQuotasTopic:
    Type: 'AWS::SNS::Topic'
    Properties:
      DisplayName: 'Update Quotas'
      KmsMasterKeyId: !Ref DefaultKey
  UpdateQuotasFunction:
    Type: 'AWS::Serverless::Function'
    Metadata:
      cfn_nag:
        rules_to_suppress:
          - id: W11
            reason: "ListServiceQuotas does not have target resource."
    Properties:
      CodeUri: functions/update_quotas/
      Handler: app.lambda_handler
      Runtime: python3.9
      Architectures:
        - x86_64
      Timeout: 60
      ReservedConcurrentExecutions: 10
      MemorySize: 512
      Events:
        UpdateQuotasSNS:
          Type: 'SNS'
          Properties:
            Topic: !GetAtt UpdateQuotasTopic.TopicArn
      VpcConfig:
        SecurityGroupIds: !Ref SecurityGroupIds
        SubnetIds: !Ref SubnetIds
      Policies:
        - AWSLambdaExecute
        - Version: '2012-10-17'
          Statement:
            - Effect: Allow
              Action:
                - 'servicequotas:ListServiceQuotaIncreaseRequestsInTemplate'
                - 'servicequotas:GetServiceQuota'
                - 'servicequotas:AssociateServiceQuotaTemplate'
                - 'iam:CreateServiceLinkedRole'
              Resource: '*'
        - Version: '2012-10-17'
          Statement:
            - Effect: Allow
              Action:
                - 'servicequotas:PutServiceQuotaIncreaseRequestIntoTemplate'
                - 'servicequotas:RequestServiceQuotaIncrease'
              Resource: !Sub 'arn:aws:servicequotas:${AWS::Region}:${AWS::AccountId}:sagemaker/*'