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

# The following template is designed to provision and configure a secure environment for data science.
# This template creates an AWS VPC, a KMS CMK, an administrator and data scientist role, and an S3 bucket.
# The template also provisions a Service Catalog portfolio and product to create notebooks into the VPC.
# Lastly the template stores outputs into Parameter Store so they can be referenced later by SC products.
Description: Data Science Environment S3 data storage

Parameters:
  SharedServiceStackSetName:
    Type: String
    Description: Common root name used across shared service cloudformation resources

  TeamName:
    Type: String
    AllowedPattern: '[A-Za-z0-9\-]*'
    Description: Please specify your team name.  Used as a suffix for team's resource names.

  EnvType:
    Description: System Environment
    Type: String
    Default: dev

Outputs:
  KMSCMK:
    Description: KMS Key ARN for the data and model buckets
    Value: !GetAtt KMSCMK.Arn
    Export:
      Name: !Sub 'ds-kms-cmk-${TeamName}-${EnvType}-arn'

  DataBucket:
    Description: Data bucket name
    Value: !Ref DataBucket
    Export:
      Name: !Sub 'ds-s3-data-${TeamName}-${EnvType}'

  ModelArtifactsBucket:
    Description: Model artifacts bucket
    Value: !Ref ModelArtifactsBucket
    Export:
      Name: !Sub 'ds-s3-models-${TeamName}-${EnvType}'

  CodeCommitUrl:
    Description: Code Commit Repository
    Value: !GetAtt CodeCommitRepo.CloneUrlHttp
    Export:
      Name: !Sub 'ds-source-${TeamName}-${EnvType}-url'

Resources:
  KMSCMK:
    Type: 'AWS::KMS::Key'
    Properties:
      Description: KMS key for S3 buckets
      EnableKeyRotation: true
      KeyPolicy:
        Id: key-policy-1
        Version: 2012-10-17
        Statement:
          - Sid: Enable IAM User Permissions
            Effect: Allow
            Principal:
              AWS: !Sub "arn:aws:iam::${AWS::AccountId}:root"
            Action: 'kms:*'
            Resource: '*'
          - Sid: Allow access for Key Administrators
            Effect: Allow
            Principal:
              AWS: 
                Fn::ImportValue:
                  !Sub "ds-admin-role-${TeamName}-${EnvType}-arn"
            Action:
              - 'kms:Create*'
              - 'kms:Describe*'
              - 'kms:Enable*'
              - 'kms:List*'
              - 'kms:Put*'
              - 'kms:Update*'
              - 'kms:Revoke*'
              - 'kms:Disable*'
              - 'kms:Get*'
              - 'kms:Delete*'
              - 'kms:TagResource'
              - 'kms:UntagResource'
              - 'kms:ScheduleKeyDeletion'
              - 'kms:CancelKeyDeletion'
            Resource: '*'
          - Sid: Allow access for Key Users
            Effect: Allow
            Principal:
              AWS: !Sub "arn:aws:iam::${AWS::AccountId}:root"
              # TODO: reduce to team specific role
#              AWS: !Sub "arn:aws:iam::${AWS::AccountId}:role/service-role/ds-notebook-role-${TeamName}-${EnvType}"
#              AWS: !Sub "arn:aws:iam::${AWS::AccountId}:role/ds-user-role-${TeamName}-${EnvType}"
            Action:
              - 'kms:Encrypt'
              - 'kms:Decrypt'
              - 'kms:CreateGrant'
              - 'kms:ReEncrypt'
              - 'kms:GenerateDataKey'
              - 'kms:DescribeKey'
            Resource: '*'
            Condition:
              StringNotEquals:
                'aws:SourceVpce':
                  Fn::ImportValue: 
                    !Sub "ds-s3-endpoint-${SharedServiceStackSetName}-id"
      Tags:
        - Key: TeamName
          Value: !Ref TeamName
        - Key: EnvironmentType
          Value: !Ref EnvType

  KMSCMKAlias:
    Type: 'AWS::KMS::Alias'
    Properties:
      AliasName: !Sub "alias/ds-s3-kms-cmk-${TeamName}-${EnvType}"
      TargetKeyId: !Ref KMSCMK

  KMSCMKArn:
    Type: 'AWS::SSM::Parameter'
    Properties:
      Name: !Sub "ds-kms-cmk-${TeamName}-${EnvType}-arn"
      Type: String
      Value: !GetAtt 
        - KMSCMK
        - Arn
      Description: SageMakerExecRole ARN

  DataBucket:
    Type: 'AWS::S3::Bucket'
    Properties:
      BucketName:
        !Join
          - ''
          -
            - 'ds-data-bucket-'
            - !Ref TeamName
            - '-'
            - !Ref EnvType
            - '-'
            - !Select
              - 4
              - !Split
                - '-'
                - !Select
                  - 2
                  - !Split
                    - /
                    - !Ref 'AWS::StackId'
      PublicAccessBlockConfiguration:
        BlockPublicAcls: TRUE
        BlockPublicPolicy: TRUE
        IgnorePublicAcls: TRUE
        RestrictPublicBuckets: TRUE
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: 'aws:kms'
              KMSMasterKeyID: !Ref KMSCMK
      Tags:
        - Key: TeamName
          Value: !Ref TeamName
        - Key: EnvironmentType
          Value: !Ref EnvType

  DataBucketPolicy:
    Type: 'AWS::S3::BucketPolicy'
    Properties:
      Bucket: !Ref DataBucket
      PolicyDocument:
        Statement:
          - Action:
              - 's3:GetObject'
              - 's3:PutObject'
              - 's3:ListBucket'
            Effect: Deny
            Resource:
              - !Sub "arn:aws:s3:::${DataBucket}/*"
              - !Sub "arn:aws:s3:::${DataBucket}"
            Principal: '*'
            Condition:
              StringNotEquals:
                'aws:SourceVpce':
                  Fn::ImportValue: 
                    !Sub "ds-s3-endpoint-${SharedServiceStackSetName}-id"

  ModelArtifactsBucket:
    Type: 'AWS::S3::Bucket'
    Properties:
      BucketName:
        !Join
          - ""
          -
            - "ds-model-bucket-"
            - !Ref TeamName
            - "-"
            - !Ref EnvType
            - "-"
            - !Select
              - 4
              - !Split
                - '-'
                - !Select
                  - 2
                  - !Split
                    - /
                    - !Ref 'AWS::StackId'
      PublicAccessBlockConfiguration:
        BlockPublicAcls: TRUE
        BlockPublicPolicy: TRUE
        IgnorePublicAcls: TRUE
        RestrictPublicBuckets: TRUE
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: 'aws:kms'
              KMSMasterKeyID: !Ref KMSCMK
      Tags:
        - Key: TeamName
          Value: !Ref TeamName
        - Key: EnvironmentType
          Value: !Ref EnvType

  ModelArtifactsBucketPolicy:
    Type: 'AWS::S3::BucketPolicy'
    Properties:
      Bucket: !Ref ModelArtifactsBucket
      PolicyDocument:
        Statement:
          - Action:
              - 's3:GetObject'
              - 's3:PutObject'
              - 's3:ListBucket'
            Effect: Deny
            Resource:
              - !Sub "arn:aws:s3:::${ModelArtifactsBucket}/*"
              - !Sub "arn:aws:s3:::${ModelArtifactsBucket}"
            Principal: '*'
            Condition:
              StringNotEquals:
                'aws:SourceVpce':
                  Fn::ImportValue: 
                    !Sub "ds-s3-endpoint-${SharedServiceStackSetName}-id"

  CodeCommitRepo:
    Type: 'AWS::CodeCommit::Repository'
    Properties:
      RepositoryName: !Sub 'ds-source-${TeamName}-${EnvType}'
      RepositoryDescription: Data science project code repository
      Code:
        S3: 
          Bucket: '< S3_CFN_STAGING_BUCKET >'
          Key: '< S3_CFN_STAGING_PATH >/project_template.zip'
      Tags:
        - Key: TeamName
          Value: !Ref TeamName
        - Key: EnvironmentType
          Value: !Ref EnvType

  S3DataBucketNameSSMParameter:
    Type: 'AWS::SSM::Parameter'
    Properties:
      Name: !Sub "ds-s3-data-bucket-${TeamName}-${EnvType}"
      Type: String
      Value: !Ref DataBucket
      Description: Data Science S3 data bucket name for the project

  S3ModelBucketNameSSMParameter:
    Type: 'AWS::SSM::Parameter'
    Properties:
      Name: !Sub "ds-s3-model-artifact-bucket-${TeamName}-${EnvType}"
      Type: String
      Value: !Ref ModelArtifactsBucket
      Description: Data Science S3 data bucket name for the project