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

AWSTemplateFormatVersion: '2010-09-09'
Description: |
  Create Amazon S3 buckets for Data Science Environment

Outputs:
  DataBucketName:
    Description: Name of the Amazon S3 bucket for data
    Value: !Ref DataBucketName
    Export:
      Name: !Sub ${EnvName}-${EnvType}-data-bucket-name

  ModelBucketName:
    Description: Name of the Amazon S3 bucket for models
    Value: !Ref ModelBucketName
    Export:
      Name: !Sub ${EnvName}-${EnvType}-model-bucket-name

Parameters:
  EnvName:
    Type: String
    AllowedPattern: '[a-z0-9\-]*'
    Description: Please specify your data science environment name.  Used as a suffix for environment resource names.

  EnvType:
    Description: System Environment (e.g. dev, test, prod). Used as a suffix for environment resource names.
    Type: String
    Default: dev

  DataBucketName:
    Description: S3 bucket name to store data
    Type: String

  ModelBucketName:
    Description: S3 bucket name to store models
    Type: String

  VPCEndpointS3Id:
    Description: Id of the S3 VPC endpoint
    Type: String

  S3BucketKMSKeyId:
    Description: KMS Key Id for the data and model buckets
    Type: String

Resources:

  DataBucket:
    Type: AWS::S3::Bucket
    Properties:
      BucketName: !Ref DataBucketName
      AccessControl: Private
      PublicAccessBlockConfiguration:
        BlockPublicAcls: TRUE
        BlockPublicPolicy: TRUE
        IgnorePublicAcls: TRUE
        RestrictPublicBuckets: TRUE
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: 'aws:kms'
              KMSMasterKeyID: !Ref S3BucketKMSKeyId
      Tags:
        - Key: EnvironmentName
          Value: !Ref EnvName
        - Key: EnvironmentType
          Value: !Ref EnvType

  ModelBucket:
    Type: AWS::S3::Bucket
    Properties:
      BucketName: !Ref ModelBucketName
      AccessControl: Private
      PublicAccessBlockConfiguration:
        BlockPublicAcls: TRUE
        BlockPublicPolicy: TRUE
        IgnorePublicAcls: TRUE
        RestrictPublicBuckets: TRUE
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: 'aws:kms'
              KMSMasterKeyID: !Ref S3BucketKMSKeyId
      Tags:
        - Key: EnvironmentName
          Value: !Ref EnvName
        - 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': !Ref VPCEndpointS3Id #Allow access only via S3 VPC endpoint
                'aws:PrincipalTag/Principal': 
                  #Exception for SageMaker pipeline execution role to read evaluation.json using JsonGet
                  - !Sub '${EnvName}-${EnvType}-sm-pipeline-role' 
                  #Exception for Feature Store replication to S3 bucket - you might create a dedicated FeatureStore role
                  - !Sub '${EnvName}-${EnvType}-sm-execution-role' 
              Bool:
                #Exception for - #### Check if this exception is needed ####
                'aws:ViaAWSService': 'false' 

  ModelBucketPolicy:
    Type: 'AWS::S3::BucketPolicy'
    Properties:
      Bucket: !Ref ModelBucket
      PolicyDocument:
        Statement:
          - Sid: AllowCrossAccount
            Action:
              - 's3:GetObject'
              - 's3:ListBucket'
            Effect: Allow
            Principal:
              AWS:
                - !Ref AWS::AccountId
            Resource:
              - !Sub "arn:aws:s3:::${ModelBucket}/*"
              - !Sub "arn:aws:s3:::${ModelBucket}"
          - Sid: DenyNoVPC
            Action:
              - 's3:GetObject'
              - 's3:PutObject'
              - 's3:ListBucket'
              - 's3:GetBucketAcl'
              - 's3:GetObjectAcl'
              - 's3:PutBucketAcl'
              - 's3:PutObjectAcl'
            Effect: Deny
            Resource:
              - !Sub "arn:aws:s3:::${ModelBucket}/*"
              - !Sub "arn:aws:s3:::${ModelBucket}"
            Principal: '*'
            Condition:
              StringNotEquals:
                'aws:sourceVpce': !Ref VPCEndpointS3Id
                'aws:PrincipalTag/Principal': !Sub '${EnvName}-${EnvType}-sm-pipeline-role' #Exception for SageMaker pipeline execution role
              Bool:
                #Exception for SageMaker execution role (for RegisterModel step)
                'aws:ViaAWSService': 'false'