AWSTemplateFormatVersion: 2010-09-09
Description: Template to provision OpenSearch cluster and SageMaker Notebook for Korean OpenSearch Workshop

Metadata:
  AWS::CloudFormation::Interface:
    ParameterGroups:
      - Label:
          default: Required Parameters
        Parameters:
          - AppName          
          - OpenSearchUsername
          - OpenSearchPassword
          - OpenSearchIndexName
          - SageMakerNotebookName
    ParameterLabels:      
      AppName:
        default: Name of the overall application
      OpenSearchUsername:
        default: OpenSearch cluster username
      OpenSearchPassword:
        default: OpenSearch cluster password
      OpenSearchIndexName:
        default: OpenSearch index name
      SageMakerNotebookName:
        default: Name of SageMaker Notebook Instance
      
Parameters:
  InstanceType:
    Type: String
    Description: Choosing c5.large
    Default: c5.large   
  OpenSearchUsername:
    AllowedPattern: '^[a-zA-Z0-9]+$'
    Default: master
    Description: User name for the account that will be added to the OpenSearch cluster.
    MaxLength: '25'
    MinLength: '5'
    Type: String
  OpenSearchPassword:
    AllowedPattern: '(?=^.{8,32}$)((?=.*\d)(?=.*[A-Z])(?=.*[a-z])|(?=.*\d)(?=.*[^A-Za-z0-9])(?=.*[a-z])|(?=.*[^A-Za-z0-9])(?=.*[A-Z])(?=.*[a-z])|(?=.*\d)(?=.*[A-Z])(?=.*[^A-Za-z0-9]))^.*'
    Description: Password for the account named above. Must be at least 8 characters containing letters, numbers and symbols
    MaxLength: '32'
    MinLength: '8'
    NoEcho: 'true'
    Type: String
    Default: Passw0rd!
  AppName:
    Default: opensearch-gen-ai
    Type: String
    AllowedValues: [opensearch-gen-ai]
    Description: Name of the overall application, this is used while creating the ML model endpoint.
  SageMakerNotebookName:
    Default: opensearch-workshop
    Type: String
    Description: Enter name of SageMaker Notebook instance. The notebook name must _not_ already exist in your AWS account/region.
    MinLength: 1
    MaxLength: 63
  SageMakerIAMRole:
    Description: Name of IAM role that will be created by this cloud formation template. The role name must _not_ already exist in your AWS account.
    Type: String
    Default: "LLMAppsIAMRole"  
  OpenSearchEngineVersion:
    Description: Engine Version of OpenSearch Domain
    Type: String
    Default: 'OpenSearch_2.13'

Mappings: 
  RegionMap: 
    us-west-2:
      HVM64: ami-05e491ac3791032d1
    us-east-1:
      HVM64: ami-07d07d65c47e5aa90
    eu-west-1:
      HVM64: ami-0b9fd8b55a6e3c9d5
 
Resources:

  OpenSearchSecret:
    Type: AWS::SecretsManager::Secret
    Properties:
      Name: !Sub OpenSearchSecret-${AWS::StackName}
      Description: OpenSearch username and password
      SecretString: !Sub '{ "username" : "${OpenSearchUsername}", "password" : "${OpenSearchPassword}" }'

  CodeRepository:
    Type: AWS::SageMaker::CodeRepository
    Properties:
      GitConfig:
        RepositoryUrl: https://github.com/atheanchu/opensearch-with-bedrock-workshop-kr

  NotebookInstance:
    Type: AWS::SageMaker::NotebookInstance
    Properties:
      NotebookInstanceName: opensearch-workshop-nb
      InstanceType: ml.m5d.2xlarge
      RoleArn: !GetAtt NBRole.Arn
      DefaultCodeRepository: !GetAtt CodeRepository.CodeRepositoryName

  NBRole:
    Type: AWS::IAM::Role
    Properties:
      Policies:
        - PolicyName: CustomNotebookAccess
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - "es:ESHttp*"
                Resource:
                  - !Sub arn:aws:es:${AWS::Region}:${AWS::AccountId}:domain/*
              - Effect: Allow
                Action:
                  - "s3:*"
                  - "dynamodb:*"
                Resource:
                  - "*"
              - Effect: Allow
                Action:
                  - "secretsmanager:GetSecretValue"
                Resource:
                  - !Sub arn:aws:secretsmanager:${AWS::Region}:${AWS::AccountId}:secret:* 
              - Effect: Allow
                Action:
                  - "bedrock:*"
                Resource: 
                  - "*"   
        - PolicyName: lambda-ssm-full-access 
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - "lambda:*"
                Resource:
                  - "*"
              - Effect: Allow
                Action:
                  - "ssm:*"
                Resource:
                  - "*"
        - PolicyName: AmazonOpenSearchServerlessFullAccess
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - "aoss:*"
                Resource:
                  - "*"
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AWSCloudFormationReadOnlyAccess
        - arn:aws:iam::aws:policy/TranslateReadOnly
        - arn:aws:iam::aws:policy/AmazonEC2FullAccess
        - arn:aws:iam::aws:policy/AmazonOpenSearchServiceFullAccess
        - arn:aws:iam::aws:policy/AmazonSageMakerFullAccess
        - arn:aws:iam::aws:policy/AWSLambda_FullAccess
        - arn:aws:iam::aws:policy/AWSLambdaExecute
        - arn:aws:iam::aws:policy/IAMFullAccess
        - arn:aws:iam::aws:policy/AmazonDynamoDBFullAccess
        
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
              - sagemaker.amazonaws.com
              - es.amazonaws.com
            Action:
              - 'sts:AssumeRole'

  Role:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Ref SageMakerIAMRole
      Policies:
        - PolicyName: CustomNotebookAccess
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Sid: ReadFromOpenSearch
                Effect: Allow
                Action:
                  - "es:ESHttp*"
                Resource:
                  - !Sub arn:aws:es:${AWS::Region}:${AWS::AccountId}:domain/*
              - Sid: ReadSecretFromSecretsManager
                Effect: Allow
                Action:
                  - "secretsmanager:GetSecretValue"
                Resource: !Sub "arn:aws:secretsmanager:${AWS::Region}:${AWS::AccountId}:secret:*"
              - Sid: ReadWriteFromECR
                Effect: Allow
                Action:
                  - "ecr:BatchGetImage"
                  - "ecr:BatchCheckLayerAvailability"
                  - "ecr:CompleteLayerUpload"
                  - "ecr:DescribeImages"
                  - "ecr:DescribeRepositories"
                  - "ecr:GetDownloadUrlForLayer"
                  - "ecr:InitiateLayerUpload"
                  - "ecr:ListImages"
                  - "ecr:PutImage"
                  - "ecr:UploadLayerPart"
                  - "ecr:CreateRepository"
                  - "ecr:GetAuthorizationToken"
                  - "ec2:DescribeAvailabilityZones"
                Resource: "*"
              - Sid: SageMakerAccess
                Effect: Allow
                Action:
                  - "sagemaker:*"
                Resource: "*"
              - Sid: S3Access
                Effect: Allow
                Action:
                  - "s3:*"
                Resource: "*"
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AmazonSageMakerFullAccess
        - arn:aws:iam::aws:policy/AWSCloudFormationReadOnlyAccess
        - arn:aws:iam::aws:policy/TranslateReadOnly
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
              - sagemaker.amazonaws.com
            Action:
              - 'sts:AssumeRole'

  OpenSearchServiceDomain:
    Type: AWS::OpenSearchService::Domain
    Properties:
      DomainName: 'opensearch-workshop'
      AccessPolicies:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              AWS: '*'
            Action: 'es:*'
            Resource: !Sub arn:aws:es:${AWS::Region}:${AWS::AccountId}:domain/*/*
      EngineVersion: !Sub ${OpenSearchEngineVersion}
      ClusterConfig:
        InstanceType: "r6g.2xlarge.search"
      EBSOptions:
        EBSEnabled: True
        VolumeSize: 20
        VolumeType: 'gp3'
      AdvancedSecurityOptions:
        AnonymousAuthEnabled: False
        Enabled: True
        InternalUserDatabaseEnabled: True
        MasterUserOptions:
          MasterUserName: !Sub ${OpenSearchUsername}
          MasterUserPassword: !Sub ${OpenSearchPassword}
      NodeToNodeEncryptionOptions:
        Enabled: True
      EncryptionAtRestOptions:
        Enabled: True
        KmsKeyId: alias/aws/es
      DomainEndpointOptions:
        EnforceHTTPS: True
  NoriPackageAssociationLambdaFunction:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |
          import boto3
          import cfnresponse

          nori_pkg_id = {}
          nori_pkg_id['us-east-1'] = {
              '2.3': 'G196105221',
              '2.5': 'G240285063',
              '2.7': 'G16029449', 
              '2.9': 'G60209291',
              '2.11': 'G181660338',
              '2.13': 'G225840180'
          }

          nori_pkg_id['us-west-2'] = {
              '2.3': 'G94047474',
              '2.5': 'G138227316',
              '2.7': 'G182407158', 
              '2.9': 'G226587000',
              '2.11': 'G79602591',
              '2.13': 'G123782433'
          }

          def lambda_handler(event, context):
              print(event)
              print('context: {}'.format(context))
              responseStatus = cfnresponse.SUCCESS
              responseData = {}
              
              if event['RequestType'] == 'Create':
                  try:
                      region = event['ResourceProperties']['Region']
                      domain_name = event['ResourceProperties']['DomainName']
                      version = event['ResourceProperties']['EngineVersion'].split('_')[1]
                      
                      opensearch = boto3.client('opensearch')
                      pkg_response = opensearch.associate_package(
                          PackageID=nori_pkg_id[region][version],
                          DomainName=domain_name
                      )
                      responseData['PackageAssociationMessage'] = 'Package Installed'
                  except Exception as e:
                      responseStatus = cfnresponse.FAILED
                      responseData = {'Error': str(e)}
              
              cfnresponse.send(event, context, responseStatus, responseData, 'CustomResourceNori')

      Handler: index.lambda_handler
      Runtime: python3.9
      Role: !GetAtt NoriPackageAssociationLambdaRole.Arn

  NoriPackageAssociationLambdaRole:
    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
        - arn:aws:iam::aws:policy/AmazonOpenSearchServiceFullAccess

  NoriPackageAssociation:
    Type: Custom::NoriPackageAssociation
    Properties:
      ServiceToken: !GetAtt NoriPackageAssociationLambdaFunction.Arn
      Region: !Ref AWS::Region
      DomainName: !Ref OpenSearchServiceDomain
      EngineVersion: !Ref OpenSearchEngineVersion

  s3BucketKnowledgeBase:
    Type: AWS::S3::Bucket
    Properties:
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: "AES256"
      VersioningConfiguration:
        Status: Enabled


Outputs:
  OpenSearchDomainEndpoint:
    Description: OpenSearch domain endpoint
    Value:
      'Fn::GetAtt':
        - OpenSearchServiceDomain
        - DomainEndpoint

  OpenSourceDomainArn:
    Description: OpenSearch domain ARN
    Value:
      'Fn::GetAtt':
        - OpenSearchServiceDomain
        - Arn

  OpenSearchDomainName:
    Description: OpenSearch domain name
    Value: !Ref OpenSearchServiceDomain

  Region:
    Description: Deployed Region
    Value: !Ref AWS::Region

  SageMakerNotebookURL:
    Description: SageMaker Notebook Instance
    Value: !Join
      - ''
      - - !Sub 'https://console.aws.amazon.com/sagemaker/home?region=${AWS::Region}#/notebook-instances/openNotebook/'
        - !GetAtt NotebookInstance.NotebookInstanceName
        - '?view=classic'
        
  OpenSearchSecret:
    Description: Name of the OpenSearch secret in Secrets Manager
    Value: !Ref OpenSearchSecret

  s3BucketKnowledgeBase:
    Description: S3 bucket name for Bedrock Agent Knowledge Base
    Value: !Ref s3BucketKnowledgeBase