# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

---
AWSTemplateFormatVersion: "2010-09-09"
Description: This stack deploys the infrastructure used in the blog post, "Explore image analysis results from Amazon Rekognition using Amazon SageMaker and Amazon DocumentDB".

Parameters:

  DocDBUsername:
    Type: String
    Description: Username for the Amazon DocumentDB cluster 

  DocDBPassword:
    Type: String
    Description: Password for the Amazon DocumentDB cluster
    NoEcho: true
    MinLength: 8

  DocDBInstanceType:
    Type: String
    Description: Instance type for DocumentDB cluster
    Default: db.r5.large
    AllowedValues:
      - db.t3.medium
      - db.r5.large
      - db.r5.xlarge
      - db.r5.2xlarge
      - db.r5.4xlarge
      - db.r5.12xlarge
      - db.r5.24xlarge

  SageMakerInstanceType:
    Type: String
    Description: Instance type for SageMaker
    Default: ml.t3.xlarge
    AllowedValues:
      - ml.c5.xlarge
      - ml.m5.xlarge
      - ml.t3.2xlarge
      - ml.t3.xlarge

  S3BucketName:
    Type: String
    Description: Bucket to store images for analysis


Metadata:
  AWS::CloudFormation::Interface: 
    ParameterGroups: 
      - 
        Label: 
          default: "DocumentDB Configuration"
        Parameters: 
          - DocDBInstanceType
          - DocDBUsername
          - DocDBPassword
    ParameterLabels: 
      DocDBInstanceType: 
        default: "Instance type for DocumentDB cluster"
      DocDBUsername: 
        default: "DocumentDB master username"
      DocDBPassword: 
        default: "DocumentDB master password"
        
Mappings:
  SubnetConfig:
    VPC:
      CIDR: '10.0.0.0/16'
    PublicOne:
      CIDR: '10.0.0.0/24'
    PrivateOne:
      CIDR: '10.0.100.0/24'
    PrivateTwo:
      CIDR: '10.0.101.0/24'
    PrivateThree:
      CIDR: '10.0.102.0/24'

Resources:
# Networking
  VPC:
    Type: AWS::EC2::VPC
    Properties:
      EnableDnsSupport: true
      EnableDnsHostnames: true
      CidrBlock: !FindInMap ['SubnetConfig', 'VPC', 'CIDR']
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-VPC

  PublicSubnetOne:
    Type: AWS::EC2::Subnet
    Properties:
      AvailabilityZone:
         Fn::Select:
         - 0
         - Fn::GetAZs: {Ref: 'AWS::Region'}
      VpcId: !Ref VPC
      CidrBlock: !FindInMap ['SubnetConfig', 'PublicOne', 'CIDR']
      MapPublicIpOnLaunch: true
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-PublicOne

  PrivateSubnetOne:
    Type: AWS::EC2::Subnet
    Properties:
      AvailabilityZone:
         Fn::Select:
         - 0
         - Fn::GetAZs: {Ref: 'AWS::Region'}
      VpcId: !Ref VPC
      CidrBlock: !FindInMap ['SubnetConfig', 'PrivateOne', 'CIDR']
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-PrivateOne

  PrivateSubnetTwo:
    Type: AWS::EC2::Subnet
    Properties:
      AvailabilityZone:
         Fn::Select:
         - 1
         - Fn::GetAZs: {Ref: 'AWS::Region'}
      VpcId: !Ref VPC
      CidrBlock: !FindInMap ['SubnetConfig', 'PrivateTwo', 'CIDR']
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-PrivateTwo

  PrivateSubnetThree:
    Type: AWS::EC2::Subnet
    Properties:
      AvailabilityZone:
         Fn::Select:
         - 2
         - Fn::GetAZs: {Ref: 'AWS::Region'}
      VpcId: !Ref VPC
      CidrBlock: !FindInMap ['SubnetConfig', 'PrivateThree', 'CIDR']
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-PrivateThree

  InternetGateway:
    Type: AWS::EC2::InternetGateway
    Properties:
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-igw
  GatewayAttachement:
    Type: AWS::EC2::VPCGatewayAttachment
    Properties:
      VpcId: !Ref VPC
      InternetGatewayId: !Ref InternetGateway
  PublicRouteTable:
    Type: AWS::EC2::RouteTable
    Properties:
      VpcId: !Ref VPC
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-PublicRouteTable
  PublicRoute:
    Type: AWS::EC2::Route
    DependsOn: GatewayAttachement
    Properties:
      RouteTableId: !Ref PublicRouteTable
      DestinationCidrBlock: 0.0.0.0/0
      GatewayId: !Ref InternetGateway
  PublicSubnetOneRouteTableAssociation:
    Type: AWS::EC2::SubnetRouteTableAssociation
    Properties:
      SubnetId: !Ref PublicSubnetOne
      RouteTableId: !Ref PublicRouteTable

# DocumentDB
  DocumentDBSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties: 
      GroupDescription: Amazon DocumentDB Security Group
      GroupName: !Sub ${AWS::StackName}-SG-DocumentDB
      VpcId: !Ref VPC
      SecurityGroupIngress:
         - IpProtocol: tcp
           CidrIp: 10.0.0.0/16
           FromPort: 27017
           ToPort: 27017
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-SG-DocumentDB
 
  DocumentDBSubnetGroup:
    Type: AWS::DocDB::DBSubnetGroup
    Properties: 
      DBSubnetGroupDescription: Subnet Group for DocumentDB
      DBSubnetGroupName: !Sub ${AWS::StackName}-SG-DocumentDB
      SubnetIds: 
         - !Ref PrivateSubnetOne
         - !Ref PrivateSubnetTwo
         - !Ref PrivateSubnetThree
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-SG-DocumentDB

  DocumentDBCluster:  
    Type: AWS::DocDB::DBCluster
    Properties: 
      DBClusterIdentifier: !Sub ${AWS::StackName}-DocumentDB
      MasterUsername: !Ref DocDBUsername
      MasterUserPassword: !Ref DocDBPassword
      DBSubnetGroupName : !Ref DocumentDBSubnetGroup
      StorageEncrypted: yes
      Tags: 
        - Key: Name
          Value: !Sub ${AWS::StackName}-DocumentDB
      VpcSecurityGroupIds: 
        - !Ref DocumentDBSecurityGroup
    DependsOn: VPC

  DocumentDBInstanceOne: 
    Type: AWS::DocDB::DBInstance
    Properties: 
      DBClusterIdentifier: !Ref DocumentDBCluster
      DBInstanceClass: !Ref DocDBInstanceType
      Tags: 
        - Key: Name
          Value: !Sub ${AWS::StackName}-DocumentDBInstance1

  DocumentDBInstanceTwo: 
    Type: AWS::DocDB::DBInstance
    Properties: 
      DBClusterIdentifier: !Ref DocumentDBCluster
      DBInstanceClass: !Ref DocDBInstanceType
      Tags: 
        - Key: Name
          Value: !Sub ${AWS::StackName}-DocumentDBInstance2

  DocumentDBInstanceThree: 
    Type: AWS::DocDB::DBInstance
    Properties: 
      DBClusterIdentifier: !Ref DocumentDBCluster
      DBInstanceClass: !Ref DocDBInstanceType
      Tags: 
        - Key: Name
          Value: !Sub ${AWS::StackName}-DocumentDBInstance3

  DocDBSecret:
    Type: 'AWS::SecretsManager::Secret'
    Properties:
      Name: !Sub ${AWS::StackName}-DocDBSecret
      Description: This secret has the credentials for the DocumentDB cluster
      SecretString:
        !Join
          - '' 
          - - '{"username":"'
            - !Ref DocDBUsername
            - '","password":"'
            - !Ref DocDBPassword
            - '", "ssl": true}'
  
  SecretDocDBClusterAttachment:
    Type: AWS::SecretsManager::SecretTargetAttachment
    Properties:
      SecretId: !Ref DocDBSecret
      TargetId: !Ref DocumentDBCluster
      TargetType: AWS::DocDB::DBCluster

# SageMaker
  SageMakerRole:
    Type: AWS::IAM::Role
    Properties: 
      AssumeRolePolicyDocument: 
        Version: "2012-10-17"
        Statement: 
          - 
            Effect: "Allow"
            Principal: 
              Service: 
                - "sagemaker.amazonaws.com"
            Action: 
              - "sts:AssumeRole"
      Description: SageMaker notebook instance role granting read access to Secrets Manager
      RoleName: !Sub ${AWS::StackName}-SageMakerRoleName
      Tags: 
        - Key: Name
          Value: !Sub ${AWS::StackName}-SageMakerRoleName

  S3AccessPolicy:
    Type: AWS::IAM::Policy
    Properties:
      PolicyName: !Sub ${AWS::StackName}-S3AccessPolicy
      Roles: 
        - !Ref SageMakerRole
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - 
            Effect: Allow
            Action:
              - s3:GetBucketLocation
              - s3:GetObject
              - s3:ListBucket
              - s3:ListBucketMultipartUploads
              - s3:ListMultipartUploadParts
              - s3:AbortMultipartUpload
              - s3:CreateBucket
              - s3:PutObject
            Resource: 
              - !Sub "arn:aws:s3:::${S3BucketName}"
              - !Sub "arn:aws:s3:::${S3BucketName}/*"

  RekognitionPolicy:
    Type: AWS::IAM::Policy
    Properties:
      PolicyName: !Sub ${AWS::StackName}-RekognitionPolicy
      Roles: 
        - !Ref SageMakerRole
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - 
            Effect: Allow
            Action:
              - "rekognition:*"
            Resource: "*"

  DocumentDBSecretPolicy:
    Type: AWS::IAM::Policy
    Properties:
      PolicyName: !Sub ${AWS::StackName}-DocumentDBSecretPolicy
      Roles: 
        - !Ref SageMakerRole
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - 
            Effect: Allow
            Action:
              - secretsmanager:GetSecretValue
            Resource: !Ref DocDBSecret

  SageMakerNotebookInstance:
    Type: AWS::SageMaker::NotebookInstance
    Properties: 
      InstanceType: !Ref SageMakerInstanceType
      NotebookInstanceName: !Sub ${AWS::StackName}-SageMakerInstance
      RoleArn: !GetAtt SageMakerRole.Arn
      DirectInternetAccess: Enabled
      LifecycleConfigName: !GetAtt SageMakerNotebookInstanceLifecycleConfig.NotebookInstanceLifecycleConfigName
      SecurityGroupIds: 
        - !Ref DocumentDBSecurityGroup
      SubnetId: !Ref PublicSubnetOne
      Tags:
        - Key: Name
          Value: !Sub ${AWS::StackName}-SageMakerInstance
      VolumeSizeInGB: 5
    DependsOn: DocumentDBCluster

  SageMakerNotebookInstanceLifecycleConfig:
    Type: "AWS::SageMaker::NotebookInstanceLifecycleConfig"
    Properties:
      OnStart:
        - Content:
            Fn::Base64: 
              !Sub |
                #!/bin/bash
                sudo -u ec2-user -i <<'EOF'
                source /home/ec2-user/anaconda3/bin/activate python3
                pip install --upgrade pymongo
                pip install --upgrade ipyplot
                source /home/ec2-user/anaconda3/bin/deactivate
                cd /home/ec2-user/SageMaker
                wget https://s3.amazonaws.com/rds-downloads/rds-combined-ca-bundle.pem
                # wget https://github.com/aws-samples/documentdb-sagemaker-example/raw/main/rekognition/script.ipynb
                mkdir pics
                cd pics
                wget https://github.com/aws-samples/documentdb-sagemaker-example/raw/main/rekognition/pics.zip
                unzip pics.zip
                rm pics.zip
                EOF


Outputs:
  StackName:
    Value: !Sub ${AWS::StackName}
  VpcId:
    Value: !Ref VPC
  PrivateSubnetOne:
    Value: !Ref PrivateSubnetOne
  PrivateSubnetTwo:
    Value: !Ref PrivateSubnetTwo
  PrivateSubnetThree:
    Value: !Ref PrivateSubnetThree
  PublicSubnetOne:
    Value: !Ref PublicSubnetOne
  DocumentDBClusterName:
    Value: !Ref DocumentDBCluster
  DocumentDBCluster:
    Value: !GetAtt DocumentDBCluster.Endpoint
  DocDBSecret:
    Value: !Ref DocDBSecret
  SageMakerRole:
    Value: !Ref SageMakerRole
  SageMakerNotebookInstance:
    Value: !Ref SageMakerNotebookInstance
  SageMakerNotebookInstanceLifecycleConfig:
    Value: !Ref SageMakerNotebookInstanceLifecycleConfig