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

AWSTemplateFormatVersion: '2010-09-09'
Description: |
  Create the base infrastructure for SageMaker secure deployment: IAM roles and shared VPC network resources

Metadata:
  AWS::CloudFormation::Interface:
    ParameterGroups:
      - Label: 
          default: Stack name
        Parameters:
          - StackSetName
      - Label:
          default: Deployment options
        Parameters:
          - CreateIAMRoles
          - CreateSharedServices
          - CreateSCPortfolio
      - Label:
          default: IAM Role ARNs (only needed if created outside of this stack)
        Parameters:
          - DSAdministratorRoleArn
          - SecurityControlExecutionRoleArn
          - SCLaunchRoleArn
    ParameterLabels:
      StackSetName:
        default: Common stack set name for core infrastructure resources
      CreateIAMRoles:
        default: Create IAM roles
      CreateSharedServices:
        default: Create shared services resources
      CreateSCPortfolio:
        default: Make a SageMaker Studio environment available as an AWS Service Catalog product
      DSAdministratorRoleArn:
        default: Data Science Administrator role ARN
      SCLaunchRoleArn:
        default: Service Catalog launch role ARN
      SecurityControlExecutionRoleArn:
        default: Security Control execution role ARN

Outputs:
  AssumeDSAdministratorRole:
    Description: URL for assuming the role of a cross-environment data science admin
    Value: !If
      - IAMRolesCondition
      - !GetAtt IAMSharedRoles.Outputs.AssumeDSAdministratorRole
      - !Sub
        - 'https://signin.aws.amazon.com/switchrole?account=${AWS::AccountId}&roleName=${DSAdmininstratorRole}&displayName=${DSAdmininstratorRole}'
        - DSAdmininstratorRole: !Select ['1', !Split [ '/', !Select ['5', !Split [':', !Ref DSAdministratorRoleArn ]]]]
    
  SCLaunchRoleArn:
    Description: Service Catalog Launch role ARN
    Value: !If
      - IAMRolesCondition
      - !GetAtt IAMSharedRoles.Outputs.SCLaunchRoleArn
      - !Ref SCLaunchRoleArn
    Export:
      Name: 'ds-service-catalog-launch-role-arn'

  ProductId:
    Description: Service Catalog data science product Id
    Value: !If 
      - SCPortfolioCondition
      - !GetAtt SharedServiceCatalogPortfolio.Outputs.ProductId
      - ''

  PortfolioId:
    Description: Service Catalog data science portfolio Id
    Value: !If 
      - SCPortfolioCondition
      - !GetAtt SharedServiceCatalogPortfolio.Outputs.PortfolioId
      - ''

  ProductName:
    Description: Service Catalog data science product name
    Value: !If 
      - SCPortfolioCondition
      - !GetAtt SharedServiceCatalogPortfolio.Outputs.ProductName
      - ''

  ProvisioningArtifactIds:
    Description: Service Catalog data science provisioning artifact Ids
    Value: !If 
      - SCPortfolioCondition
      - !GetAtt SharedServiceCatalogPortfolio.Outputs.ProvisioningArtifactIds
      - ''

  ProvisioningArtifactNames:
    Description: Service Catalog data science provisioning artifact names
    Value: !If 
      - SCPortfolioCondition
      - !GetAtt SharedServiceCatalogPortfolio.Outputs.ProvisioningArtifactNames
      - ''

  GetEnvironmentConfigurationLambdaArn:
    Description: Lambda function to retrieve the data science environment configuration 
    Value: !GetAtt GetEnvironmentConfigurationLambda.Arn
    Export:
      Name: 'ds-get-environment-configuration-lambda-arn'

Parameters:
  StackSetName:
    Type: String 
    Description: A name to be used across core infrastructure nested stacks

  CreateIAMRoles:
    Type: String
    Default: 'YES'
    AllowedValues:
      - 'YES'
      - 'NO'
    Description: Set to NO if you do not want to provision IAM roles (you must have already these roles provisioned up-front)

  DSAdministratorRoleArn:
    Type: String 
    Default: ''
    Description: Data Science Administrator role ARN if it is created outside of this stack. Must be provided if Create IAM Roles = NO

  SCLaunchRoleArn:
    Type: String
    Default: ''
    Description: Service Catalog Launch role ARN if it is create outside of this stack. Must be provided if Create IAM Roles = NO

  SecurityControlExecutionRoleArn:
    Default: ''
    Type: String 
    Description: Execution role ARN for security controls if it is create outside of this stack. Must be provided if Create IAM Roles = NO 

  CreateSharedServices:
    Type: String
    Default: 'NO'
    AllowedValues:
        - 'YES'
        - 'NO'
    Description: Set to YES if you do want to provision the shared services VPC network and PyPi mirror repository

  CreateSCPortfolio:
    Type: String
    Default: 'YES'
    AllowedValues:
        - 'YES'
        - 'NO'
    Description: Set to YES if you want to make a SageMaker Studio environment available as a Service Catalog product

Conditions:
  IAMRolesCondition: !Equals [ !Ref CreateIAMRoles, 'YES' ]
  SharedServicesCondition: !Equals [ !Ref CreateSharedServices, 'YES' ]
  SCPortfolioCondition: !Equals [ !Ref CreateSCPortfolio, 'YES' ]

Rules:
  IAMRoles:
    RuleCondition: !Equals [ !Ref CreateIAMRoles, 'NO' ]
    Assertions:
      - Assert: !And
          - !Not [ !Equals [ !Ref DSAdministratorRoleArn, '' ] ]
          - !Not [ !Equals [ !Ref SCLaunchRoleArn, '' ] ]
          - !Not [ !Equals [ !Ref SecurityControlExecutionRoleArn, '' ] ]
        AssertDescription: ARNs must be provided for all IAM roles if Create IAM Roles = NO

Resources:

  GetEnvironmentConfigurationLambda:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |
          # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
          # SPDX-License-Identifier: MIT-0
          import json
          import boto3
          import cfnresponse

          sm = boto3.client("sagemaker")
          ssm = boto3.client("ssm")

          def get_environment(project_name, ssm_params):
              r = sm.describe_domain(
                      DomainId=sm.describe_project(
                          ProjectName=project_name
                          )["CreatedBy"]["DomainId"]
                  )
              del r["ResponseMetadata"]
              del r["CreationTime"]
              del r["LastModifiedTime"]
              r = {**r, **r["DefaultUserSettings"]}
              del r["DefaultUserSettings"]

              i = {
                  **r,
                  **{t["Key"]:t["Value"] 
                      for t in sm.list_tags(ResourceArn=r["DomainArn"])["Tags"] 
                      if t["Key"] in ["EnvironmentName", "EnvironmentType"]}
              }

              for p in ssm_params:
                  try:
                      i[p["VariableName"]] = ssm.get_parameter(Name=f"{i['EnvironmentName']}-{i['EnvironmentType']}-{p['ParameterName']}")["Parameter"]["Value"]
                  except:
                      i[p["VariableName"]] = ""

              return i
              
          def lambda_handler(event, context):
              try:
                  response_status = cfnresponse.SUCCESS
                  r = {}

                  if 'RequestType' in event and event['RequestType'] == 'Create':
                      r = get_environment(event["ResourceProperties"]["SageMakerProjectName"], event["ResourceProperties"]["SSMParams"])

                  print(r)
                  cfnresponse.send(event, context, response_status, r, '')

              except Exception as exception:
                  print(exception)
                  cfnresponse.send(event, context, cfnresponse.FAILED, {}, physicalResourceId=event.get('PhysicalResourceId'), reason=str(exception))

      Description: Get data science environment configuration for MLOps project setup
      Handler: index.lambda_handler
      MemorySize: 128
      Role: !If
          - IAMRolesCondition
          - !GetAtt IAMSCSageMakerProjectRoles.Outputs.AmazonSageMakerServiceCatalogProductsLaunchRoleArn
          - !Sub 'arn:${AWS::Partition}:iam::${AWS::AccountId}:role/service-role/AmazonSageMakerServiceCatalogProductsLaunchRole'
      Runtime: python3.8
      Timeout: 60
      Tags:
        - Key: DataScienceStackName
          Value: !Ref StackSetName

  LambdaExecutionRole:
    Type: 'AWS::IAM::Role'
    Condition: IAMRolesCondition
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Path: /
      Policies:
        - PolicyName: InlinePolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Sid: IAMPermission
                Effect: Allow
                Action:
                  - 'iam:GetRole'
                Resource:
                  - !Sub 'arn:*:iam::${AWS::AccountId}:*'
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
      Tags:
        - Key: DataScienceStackName
          Value: !Ref StackSetName

  CheckIfIAMRolesExistLambda:
    Type: AWS::Lambda::Function
    Condition: IAMRolesCondition
    Properties:
      Code:
        ZipFile: |
          import json
          import boto3
          import cfnresponse
          from botocore.exceptions import ClientError

          iam = boto3.client('iam')

          def lambda_handler(event, context):
              try:
                  response_status = cfnresponse.SUCCESS
                  r = {}

                  if 'RequestType' in event and event['RequestType'] == 'Create':
                      r = check_roles(event['ResourceProperties']['RoleNames'])

                  cfnresponse.send(event, context, response_status, r, '')

              except ClientError as exception:
                  print(exception)
                  cfnresponse.send(event, context, cfnresponse.FAILED, {}, physicalResourceId=event.get('PhysicalResourceId'), reason=str(exception))

          def check_roles(roles):
              ret = {}
              for r in roles:
                  try:
                      iam.get_role(RoleName=r)
                      ret[r] = ""
                  except iam.exceptions.NoSuchEntityException:
                      ret[r] = r

              return ret
      Description: Checks if specified IAM roles exist
      Handler: index.lambda_handler
      MemorySize: 128
      Role: !GetAtt LambdaExecutionRole.Arn
      Runtime: python3.8
      Timeout: 60
      Tags:
        - Key: DataScienceStackName
          Value: !Ref StackSetName

  CheckIfSageMakerSCProductsRolesExist:
    Type: Custom::CheckIfIAMRolesExist
    Condition: IAMRolesCondition
    Properties:
      ServiceToken: !GetAtt CheckIfIAMRolesExistLambda.Arn
      RoleNames:
        - AmazonSageMakerServiceCatalogProductsLaunchRole
        - AmazonSageMakerServiceCatalogProductsUseRole

  # SageMaker Project templates launch and use roles
  IAMSCSageMakerProjectRoles:
    Type: AWS::CloudFormation::Stack
    Condition: IAMRolesCondition
    Properties:
      Parameters: 
        AmazonSageMakerServiceCatalogProductsLaunchRoleName: 
          !GetAtt CheckIfSageMakerSCProductsRolesExist.AmazonSageMakerServiceCatalogProductsLaunchRole
        AmazonSageMakerServiceCatalogProductsUseRoleName: 
          !GetAtt CheckIfSageMakerSCProductsRolesExist.AmazonSageMakerServiceCatalogProductsUseRole
      TemplateURL: core-iam-sc-sm-projects-roles.yaml
      Tags:
        - Key: DataScienceStackName
          Value: !Ref StackSetName

  # Shared IAM roles for personas and services
  IAMSharedRoles:
    Type: AWS::CloudFormation::Stack
    Condition: IAMRolesCondition
    Properties:
      Parameters:
        DSAdministratorRoleName: !Sub '${StackSetName}-${AWS::Region}-DataScienceAdministrator'
        SageMakerDetectiveControlExecutionRoleName: !Sub '${StackSetName}-${AWS::Region}-DSSageMakerDetectiveControlRole'
        SCLaunchRoleName: !Sub '${StackSetName}-${AWS::Region}-DSServiceCatalogLaunchRole'
      TemplateURL: core-iam-shared-roles.yaml
      Tags:
        - Key: DataScienceStackName
          Value: !Ref StackSetName

  # Shared services VPC and related network resources
  SharedServicesNetwork:
    Type: AWS::CloudFormation::Stack
    Condition: SharedServicesCondition
    Properties:
      Parameters:
        StackSetName: !Ref StackSetName
      TemplateURL: core-shared-services-network.yaml
      Tags:
        - Key: DataScienceStackName
          Value: !Ref StackSetName

  # Private PyPi mirror
  SharedServicePyPIMirror:
    Type: AWS::CloudFormation::Stack
    DependsOn: SharedServicesNetwork
    Condition: SharedServicesCondition
    Properties:
      Parameters:
        StackSetName: !Ref StackSetName
      TemplateURL: core-shared-services-ecs.yaml
      Tags:
        - Key: DataScienceStackName
          Value: !Ref StackSetName

  # Service Catalog with a Data Science product  
  SharedServiceCatalogPortfolio:
    Type: AWS::CloudFormation::Stack
    Condition: SCPortfolioCondition
    Properties:
      Parameters:
        StackSetName: !Ref StackSetName
        SCPortfolioPrincipalRoleArn: !If
          - IAMRolesCondition
          - !GetAtt IAMSharedRoles.Outputs.DSAdministratorRoleArn
          - !Ref DSAdministratorRoleArn
        SCProductLaunchRoleArn: !If
          - IAMRolesCondition
          - !GetAtt IAMSharedRoles.Outputs.SCLaunchRoleArn
          - !Ref SCLaunchRoleArn
      TemplateURL: core-sc-shared-portfolio.yaml
      Tags:
        - Key: DataScienceStackName
          Value: !Ref StackSetName

  # Security controls for the Data Science environment
  DSSecurityControls:
    Type: AWS::CloudFormation::Stack
    Properties:
      Parameters:
        StackSetName: !Ref StackSetName
        SecurityControlExecutionRoleArn: !If
          - IAMRolesCondition
          - !GetAtt IAMSharedRoles.Outputs.SageMakerDetectiveControlExecutionRoleArn
          - !Ref SecurityControlExecutionRoleArn
      TemplateURL: core-security-controls.yaml
      Tags:
        - Key: DataScienceStackName
          Value: !Ref StackSetName