AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31
Description: >
  sagemaker-saml-backend

Parameters:
  EnvironmentName:
    Type: String
    Default: SageMakerSecureSSO
    AllowedPattern: '[a-zA-Z0-9]+'
  NetworkVpcCidr:
    Type: String
    Default: 10.100.0.0/16
    AllowedPattern: ^([0-9]{1,3}\.){3}[0-9]{1,3}($|\/(8|16|24|32))$
  NetworkPrivateSubnetCidr:
    Type: String
    Default: 10.100.10.0/24
    AllowedPattern: ^([0-9]{1,3}\.){3}[0-9]{1,3}($|\/(8|16|24|32))$
  NetworkPrivateSubnetSageMakerCidr:
    Type: String
    Default: 10.100.30.0/24
    AllowedPattern: ^([0-9]{1,3}\.){3}[0-9]{1,3}($|\/(8|16|24|32))$  
  NetworkPublicSubnetCidr:
    Type: String
    Default: 10.100.20.0/24
    AllowedPattern: ^([0-9]{1,3}\.){3}[0-9]{1,3}($|\/(8|16|24|32))$
  SageMakerDomainName:
    Type: String
    Default: SageMakerSecure
    AllowedPattern: '[a-zA-Z0-9]+'
  SSOUserId:
    Type: String
    AllowedPattern: ^[{]?-[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}
  KeyPairName: 
    Type: String
    AllowedPattern: '[a-zA-Z0-9\-\_]+'
  PublicIpCidr:
    Type: String
    AllowedPattern: ^([0-9]{1,3}\.){3}[0-9]{1,3}($|\/(8|16|24|32))$
Mappings: 
  RegionMap: 
    eu-west-1: 
      LinuxAMI: "ami-0d1bf5b68307103c2"
      WinAMI: "ami-07fc612d98a90981b"     
      datascience: "arn:aws:sagemaker:eu-west-1:470317259841:image/datascience-1.0"
      datawrangler: "arn:aws:sagemaker:eu-west-1:245179582081:image/sagemaker-data-wrangler-1.0"

    eu-central-1: 
      LinuxAMI: "ami-058e6df85cfc7760b"
      WinAMI: "ami-05a60358d5cda31c5"      
      datascience: "arn:aws:sagemaker:eu-central-1:936697816551:image/datascience-1.0"
      datawrangler: "arn:aws:sagemaker:eu-central-1:024640144536:image/sagemaker-data-wrangler-1.0"

    us-east-1: 
      LinuxAMI: "ami-04902260ca3d33422"
      WinAMI: "ami-064d05b4fe8515623"     
      datascience: "arn:aws:sagemaker:us-east-1:081325390199:image/datascience-1.0"
      datawrangler: "arn:aws:sagemaker:us-east-1:663277389841:image/sagemaker-data-wrangler-1.0"

Globals:
  Function:
    Timeout: 10
    MemorySize: 128

Resources:
  
######## CF Stacks ########

  VPCStack:
    Type: AWS::CloudFormation::Stack
    Properties: 
      Parameters:
        EnvironmentName : !Ref EnvironmentName
        NetworkVpc: !Ref NetworkVpcCidr
        NetworkPrivateSubnet: !Ref NetworkPrivateSubnetCidr
        NetworkPrivateSubnetSageMaker: !Ref NetworkPrivateSubnetSageMakerCidr
        NetworkPublicSubnet: !Ref NetworkPublicSubnetCidr
        PublicIp: !Ref PublicIpCidr
      TemplateURL: stacks/vpc.yaml

  IAMStack:
    Type: AWS::CloudFormation::Stack
    Properties: 
      Parameters:
        EnvironmentName : !Ref EnvironmentName
        NetworkPrivateSubnet: !Ref NetworkPrivateSubnetCidr
      TemplateURL: stacks/iam.yaml

  SageMakerStack:
    Type: AWS::CloudFormation::Stack
    Properties: 
      Parameters: 
        SageMakerDomainName : !Ref SageMakerDomainName
        PrivateSubnetVPCSageMakerStudio : !GetAtt VPCStack.Outputs.PrivateSubnetVPCSageMakerStudio
        SecurityGroupSageMakerStudio : !GetAtt VPCStack.Outputs.SecurityGroupSageMakerStudio
        VPCSageMaker : !GetAtt VPCStack.Outputs.VPCSageMaker
        SageMakerExecutionRoleArn: !GetAtt IAMStack.Outputs.SageMakerExecutionRoleArn
        EnvironmentName : !Ref EnvironmentName
        SSOUserId : !Ref SSOUserId
        DataScienceImage : !FindInMap
          - RegionMap
          - !Ref 'AWS::Region'
          - datascience       
      Tags: 
        - Key: EnvironmentName
          Value: !Ref EnvironmentName
      TemplateURL: stacks/sagemakerstudio.yaml

  EC2Stack:
    Type: AWS::CloudFormation::Stack
    Properties: 
      Parameters:
        EnvironmentName : !Ref EnvironmentName
        ImageIdLinux: !FindInMap
                    - RegionMap
                    - !Ref AWS::Region
                    - LinuxAMI
        ImageIdWindows: !FindInMap
                    - RegionMap
                    - !Ref AWS::Region
                    - WinAMI        
        KeyPairName: !Ref KeyPairName
        NetworkVpcCidr: !Ref NetworkVpcCidr
        PublicSubnetVPCSageMaker: !GetAtt VPCStack.Outputs.PublicSubnetVPCSageMaker
        PrivateSubnetVPCSageMaker: !GetAtt VPCStack.Outputs.PrivateSubnetVPCSageMaker
        SecurityGroupWindows: !GetAtt VPCStack.Outputs.SecurityGroupWindows
        SecurityGroupBastion: !GetAtt VPCStack.Outputs.SecurityGroupBastion
      TemplateURL: stacks/ec2.yaml  

######## API Gateway ########

  SageMakerStudioSSOAPI:
    Type: AWS::Serverless::Api
    Properties:
      StageName: prod
      EndpointConfiguration:
        Type: Private
        VpcEndpointIds: 
          - !GetAtt VPCStack.Outputs.APIGatewayVpcEndpoint   

######## Lambda Functions ########

  SAMLBackEndFunction:
    Type: AWS::Serverless::Function 
    Properties:
      ReservedConcurrentExecutions: 1
      CodeUri: functions/saml-backend/
      Role: !GetAtt IAMStack.Outputs.SAMLBackEndLambdaRoleArn
      Handler: index.handler
      Runtime: nodejs14.x
      Events:
        SAMLBackEnd:
          Type: Api
          Properties:
            Auth: 
              ResourcePolicy:
                IntrinsicVpceWhitelist:
                  - !GetAtt VPCStack.Outputs.APIGatewayVpcEndpoint
            Path: /saml
            Method: POST
            RestApiId: !Ref SageMakerStudioSSOAPI
      VpcConfig:
        SecurityGroupIds:
          - !GetAtt VPCStack.Outputs.SecurityGroupAPIGateway
        SubnetIds:
          - !GetAtt VPCStack.Outputs.PrivateSubnetVPCSageMaker

Outputs:

  SAMLBackEndApi:
    Description: "API Gateway endpoint URL acting as the Application ACS URL"
    Value: !Sub "https://${SageMakerStudioSSOAPI}.execute-api.${AWS::Region}.amazonaws.com/prod/saml"
  SAMLAudience:
    Description: "Application SAML audience"
    Value: !Sub "https://${SageMakerStudioSSOAPI}.execute-api.${AWS::Region}.amazonaws.com/"    
  SageMakerStudioDomainId:
    Description: "SageMaker Studio domain created"
    Value: !GetAtt SageMakerStack.Outputs.SageMakerStudioDomainId
  SageMakerWindowsPublicHost:
    Description: "Public IP Address to access Windows Bastion Host"
    Value: !GetAtt EC2Stack.Outputs.Ec2InstanceWindowsPublicIp   
  SageMakerWindowsPassword:
    Description: "Default password to loging to the EC2 Windows instances"
    Value: !Sub ${EnvironmentName}${NetworkVpcCidr}
  TunnelCommand:
    Description: "Command to initiate the SSH tunnnel in order to access the Windows EC2 on the private subnet over RDP"
    Value: !Join
            - ''
            - - 'ssh -i '
              - !Ref KeyPairName
              - '.pem '
              - '-A -N -L localhost:3389:'
              - !GetAtt EC2Stack.Outputs.Ec2InstanceWindowsPrivateIp
              - ':3389 ec2-user@'
              - !GetAtt EC2Stack.Outputs.Ec2InstanceBastionPublicIp