---
AWSTemplateFormatVersion: 2010-09-09
Description: A basic template for creating a Lambda-backed API Gateway for use as
  a custom identity provider in AWS Transfer. It authenticates against cognito for password based authentication and public keys are returned from S3 bucket if public key based authentication is used. You can modify the Lambda function code to do something different after deployment.

Metadata:
  AWS::CloudFormation::Interface:
    ParameterGroups:
      -
        Label:
          default: "Network Config"
        Parameters:
          - Prefix
          - VpcId
          - SubnetIdAZa
          - SubnetIdAZb
          - SubnetIdAZc
Parameters:
  Prefix:
    Description: |
      Prefix for all resource that are created by the stack | 
      use all small letters |
      Eg: transfer 
    Type: String
  VpcId:
    Description: VPC ID to deploy the transfer server endpoint in
    Type: AWS::EC2::VPC::Id
  SubnetIdAZa:
    Description: Subnet ID for AZ - a | This must be internal subnet only
    Type: AWS::EC2::Subnet::Id
  SubnetIdAZb:
    Description: Subnet ID for AZ - b | This must be internal subnet only
    Type: AWS::EC2::Subnet::Id
  SubnetIdAZc:
    Description: Subnet ID for AZ - c | This must be internal subnet only
    Type: AWS::EC2::Subnet::Id

  CertificateArn:
    Description: Certificate Arn | Use only needed if FTPS protocal needs to be used
    Type: String

  TransferSSMParameter:
    Description: Transfer SSM Parameter name | No need to change 
    Type: String
    Default: /default/aws/transfer

Conditions:
  CertificateArnNotPresent:
    Fn::Equals:
      - Ref: CertificateArn
      - ""

Rules:
  OnlyShowRelevantSubnets:
    Assertions:
      - Assert:
          'Fn::EachMemberEquals':
            - 'Fn::ValueOfAll':
              - 'AWS::EC2::Subnet::Id'
              - VpcId
            - Ref: VpcId
        AssertDescription: 'Only shows the respective subnets for the VPC'

Resources:
  
  # Role given by transfer server to authenticated user
  TransferS3AccessRole:
    Type: AWS::IAM::Role
    Properties:
      Description: IAM role used by Transfer to give access to S3 bucket after user is authenticated
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - transfer.amazonaws.com
            Action:
              - sts:AssumeRole
  
  UserAuthenticationLambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      Description: IAM role used for Lambda Execution Role for User Authentication Lambda
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Principal:
            Service:
            - lambda.amazonaws.com
          Action:
          - sts:AssumeRole
      ManagedPolicyArns:
        - Fn::Sub: arn:${AWS::Partition}:iam::aws:policy/service-role/AWSLambdaVPCAccessExecutionRole
  
  TransferCloudWatchLoggingRole:
    Type: AWS::IAM::Role
    Properties:
      Description: IAM role used by Transfer to log requests to CloudWatch
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - transfer.amazonaws.com
            Action:
              - sts:AssumeRole
      Path: /
      ManagedPolicyArns:
        - >-
          arn:aws:iam::aws:policy/service-role/AWSTransferLoggingAccess
  
  TransferApiCloudWatchLogsRole:
    Type: AWS::IAM::Role
    Properties:
      Description: "IAM role used by API Gateway to log API requests to CloudWatch"
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Principal:
            Service:
              - apigateway.amazonaws.com
          Action:
            - sts:AssumeRole
      Path: /
      ManagedPolicyArns:
        - >-
          arn:aws:iam::aws:policy/service-role/AmazonAPIGatewayPushToCloudWatchLogs
  
  TransferIdentityProviderRole:
    Type: AWS::IAM::Role
    Properties:
      Description: Allows transfer to invoke API
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Principal:
            Service: transfer.amazonaws.com
          Action:
            - sts:AssumeRole
  

  TransferS3Bucket:
    Type: AWS::S3::Bucket
    Properties:
      AccessControl: Private
      BucketName: 
        Fn::Sub: '${AWS::Region}-${AWS::AccountId}-${Prefix}-data'
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: AES256
      PublicAccessBlockConfiguration:
        BlockPublicAcls: true
        BlockPublicPolicy: true
        IgnorePublicAcls: true
        RestrictPublicBuckets: true
      VersioningConfiguration:
        Status: Enabled
  
  TransferPublicKeyS3Bucket:
    Type: AWS::S3::Bucket
    Properties:
      AccessControl: Private
      BucketName: 
        Fn::Sub: '${AWS::Region}-${AWS::AccountId}-${Prefix}-public-keys'
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: AES256
      PublicAccessBlockConfiguration:
        BlockPublicAcls: true
        BlockPublicPolicy: true
        IgnorePublicAcls: true
        RestrictPublicBuckets: true
      VersioningConfiguration:
        Status: Enabled


  UserPool:
    Type: AWS::Cognito::UserPool
    Properties:
      UserPoolName: 
        Fn::Sub: ${Prefix}-user-pool
      UsernameConfiguration:
        CaseSensitive: true
      AccountRecoverySetting:
        RecoveryMechanisms:
          - Name: admin_only
            Priority: 1
      Policies: 
        PasswordPolicy: 
          MinimumLength: 8
          RequireLowercase: True
          RequireNumbers: True
          RequireSymbols: True
          RequireUppercase: True
          TemporaryPasswordValidityDays: 1
      UserPoolTags:
        StackName: 
          Ref: AWS::StackName
  
  UserPoolClient:
    Type: AWS::Cognito::UserPoolClient
    Properties:
      ClientName: 
        Fn::Sub: ${Prefix}-client
      GenerateSecret: false
      UserPoolId: 
        Ref: UserPool
      ExplicitAuthFlows:
        - ALLOW_USER_PASSWORD_AUTH
        - ALLOW_REFRESH_TOKEN_AUTH
        - ALLOW_USER_SRP_AUTH

  UserAuthenticationLambda:
    Type: AWS::Lambda::Function
    Properties:
      Environment: 
        Variables: 
          transferssmparameter: 
            Fn::Join:
              - ""
              - - Ref: TransferSSMParameter
                - "/"
                - Ref: Prefix
          bucketname: 
            Ref: TransferS3Bucket
          transferaccessrole:
            Fn::GetAtt: TransferS3AccessRole.Arn
          publickkeysbucketname:
            Ref: TransferPublicKeyS3Bucket
      Code:
        ZipFile: |
          import boto3
          import json
          import os
          BUCKETNAME = os.environ['bucketname']
          PUBLICKEYSBUCKETNAME = os.environ['publickkeysbucketname']
          TRANSFERACCESSROLE = os.environ['transferaccessrole']
          PRIMARY_REGION = os.environ['AWS_REGION']
          SSM_PREFIX = os.environ['transferssmparameter']
          def lambda_handler(event, context):    
            auth_response = {}
            data_ret = {}
            PublicKeysList = []
            # Check if password based authentication or not
            if event.get("password", "") != "":
              # Logic for passwor-based authentication
              print ("Password Authentication for " + event["username"])
              cognito = boto3.client("cognito-idp", region_name=PRIMARY_REGION)
              ssm = boto3.client('ssm', region_name=PRIMARY_REGION)
              ssm_parameter_string = ssm.get_parameter (Name=SSM_PREFIX,WithDecryption=True)
              ssm_parameter_json = json.loads(ssm_parameter_string['Parameter']['Value'])
              # Initiate auth with Cognito
              try:
                auth_response = cognito.initiate_auth(
                  ClientId=ssm_parameter_json['userpoolclientid'],
                  AuthFlow="USER_PASSWORD_AUTH",
                  AuthParameters={ "USERNAME": event["username"], "PASSWORD": event["password"]})
              except cognito.exceptions.UserNotFoundException as unf_error:
                # User not found Exception
                print ("User Not Found Exception: " + event["username"])
                return {}
              except cognito.exceptions.NotAuthorizedException as na_error:
                # User not authorized exception
                print ("User not authorized: " + event["username"])
                return {}
              # Checks if User password needs to changed 
              if "ChallengeName" in auth_response:
                  if "NEW_PASSWORD_REQUIRED" == auth_response["ChallengeName"]:
                      raise Exception("Please change user password.")
              data_ret["Role"] = TRANSFERACCESSROLE
            else:
              # For Password-less authentication
              if event["protocol"] == 'SFTP' :
                # Public key-based authentication 
                print ("SFTP Authentication for " + event["username"])
                s3 = boto3.client('s3')
                # Get list of objects with prefix of user name
                result = s3.list_objects(Bucket=PUBLICKEYSBUCKETNAME, Prefix=event["username"])
                # Check if list is empty return error
                if result.get('Contents') is None:
                  print ("User Key Not Found: " + event["username"])
                  return {}
                # Get data from all the available public keys 
                for object in result.get('Contents'):
                  data = s3.get_object(Bucket=PUBLICKEYSBUCKETNAME, Key=object.get('Key'))
                  if not object.get('Key').endswith("/"):
                    PublicKeysList.append(data['Body'].read().decode())
                data_ret["PublicKeys"] = PublicKeysList
                data_ret["Role"] = TRANSFERACCESSROLE
              else: 
                print ("Password field cannot be empty")
                return {}
            directorymapping = [{"Entry": "/" + event["username"], "Target": "/"+ BUCKETNAME + "/" + event["username"]}
              ,{"Entry": "/publickeys", "Target": "/"+ PUBLICKEYSBUCKETNAME + "/" + event["username"]}]
            data_pol = {
              "Version": "2012-10-17",
              "Statement": [{
                "Sid": "AllowListAccessToBucket",
                "Action": ["s3:ListBucket"],
                "Effect": "Allow",
                "Resource": ["arn:aws:s3:::" + BUCKETNAME, "arn:aws:s3:::" + PUBLICKEYSBUCKETNAME],
                "Condition": { "StringLike": { "s3:prefix": [ event["username"] + "/*", event["username"]]}}},
                {"Sid": "GetAccessForPublicKeys",
                "Effect": "Allow",
                "Action": ["s3:GetObject","s3:GetObjectAcl","s3:GetObjectVersion","s3:GetBucketLocation"],
                "Resource": ["arn:aws:s3:::" + PUBLICKEYSBUCKETNAME + "/" + event["username"]
                    ,"arn:aws:s3:::" + PUBLICKEYSBUCKETNAME + "/" + event["username"] + "/*"]},
                {"Sid": "PublicKeysPubFilePutAccess",
                "Effect": "Allow",
                "Action": ["s3:PutObject","s3:PutObjectAcl"],
                "Resource": ["arn:aws:s3:::" + PUBLICKEYSBUCKETNAME + "/" + event["username"] + "/*.pub"]},
                {"Sid": "TransferDataBucketAccess",
                "Effect": "Allow",
                "Action": ["s3:PutObject","s3:PutObjectAcl","s3:GetObject","s3:GetObjectAcl","s3:GetObjectVersion","s3:GetBucketLocation","s3:DeleteObject","s3:DeleteObjectVersion"],
                "Resource": ["arn:aws:s3:::" + BUCKETNAME + "/" + event["username"],
                    "arn:aws:s3:::" + BUCKETNAME + "/" + event["username"] + "/*"]}]}
            data_ret["Policy"] = json.dumps(data_pol)
            data_ret["HomeDirectoryType"] = "LOGICAL"
            data_ret["HomeDirectoryDetails"] = json.dumps(directorymapping)
            return data_ret
      Description: A function to provide IAM roles and policies for given user and serverId.
      Handler: index.lambda_handler
      Role: 
        Fn::GetAtt: UserAuthenticationLambdaExecutionRole.Arn
      Runtime: python3.11

  
  CognitoParameter:
    Type: AWS::SSM::Parameter
    Properties:
      Name: 
        Fn::Join:
          - ""
          - - Ref: TransferSSMParameter
            - "/"
            - Ref: Prefix
      Type: String
      Description: Cognito User Pool Client ID
      Value: 
        Fn::Join:
          - ""
          - - "{"
            - Fn::Join:
                - ""
                - - '"userpoolid":"'
                  - Ref: UserPool
                  - '",'
            - Fn::Join:
                - ""
                - - '"userpoolclientid":"'
                  - Ref: UserPoolClient
                  - '"'
            - "}"
  
  TransferSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: transfer-server-security-group
      VpcId: 
        Ref: VpcId
      SecurityGroupIngress:
        - IpProtocol: tcp
          FromPort: 21
          ToPort: 22
          CidrIp: 0.0.0.0/0
        - IpProtocol: tcp
          FromPort: 8192
          ToPort: 8200
          CidrIp: 0.0.0.0/0
      SecurityGroupEgress:
        - CidrIp: 0.0.0.0/0
          IpProtocol: "-1"

  TransferServer:
    Type: AWS::Transfer::Server
    Properties:
      EndpointType: VPC
      Domain: S3
      Certificate: 
        Fn::If:
          - CertificateArnNotPresent
          - Ref: AWS::NoValue
          - Ref: CertificateArn
      Protocols:
        - FTP
        - SFTP
        - Fn::If:
            - CertificateArnNotPresent
            - Ref: AWS::NoValue
            - FTPS
      EndpointDetails:
        SubnetIds:
          - Ref: SubnetIdAZa
          - Ref: SubnetIdAZb
          - Ref: SubnetIdAZc
        VpcId:
          Ref: VpcId
        SecurityGroupIds:
          - Ref: TransferSecurityGroup
      IdentityProviderDetails:
        InvocationRole:
          Fn::GetAtt: TransferIdentityProviderRole.Arn
        Url:
          Fn::Join:
            - ''
            - - https://
              - Ref: CustomIdentityProviderApi
              - .execute-api.
              - Ref: AWS::Region
              - '.'
              - Ref: AWS::URLSuffix
              - '/'
              - Ref: ApiStage
      IdentityProviderType: API_GATEWAY
      LoggingRole:
        Fn::GetAtt: TransferCloudWatchLoggingRole.Arn
  

  CustomIdentityProviderApi:
    Type: AWS::ApiGateway::RestApi
    Properties:
      Name: Transfer Custom Identity Provider API
      Description: API used for Authentication requests
      FailOnWarnings: true
      EndpointConfiguration:
        Types:
          - REGIONAL
  ApiLoggingAccount:
    Type: AWS::ApiGateway::Account
    DependsOn:
    - CustomIdentityProviderApi
    Properties:
      CloudWatchRoleArn:
        Fn::GetAtt: TransferApiCloudWatchLogsRole.Arn
  ApiStage:
    Type: AWS::ApiGateway::Stage
    Properties:
      DeploymentId:
        Ref: ApiDeployment
      MethodSettings:
      - DataTraceEnabled: false
        HttpMethod: "*"
        LoggingLevel: INFO
        ResourcePath: "/*"
      RestApiId:
        Ref: CustomIdentityProviderApi
      StageName: 
        Fn::Sub: ${Prefix}
  ApiDeployment:
    DependsOn:
    - GetUserConfigRequest
    Type: AWS::ApiGateway::Deployment
    Properties:
      RestApiId:
        Ref: CustomIdentityProviderApi
  ServersResource:
    Type: AWS::ApiGateway::Resource
    Properties:
      RestApiId:
        Ref: CustomIdentityProviderApi
      ParentId:
        Fn::GetAtt:
        - CustomIdentityProviderApi
        - RootResourceId
      PathPart: servers
  ServerIdResource:
    Type: AWS::ApiGateway::Resource
    Properties:
      RestApiId:
        Ref: CustomIdentityProviderApi
      ParentId:
        Ref: ServersResource
      PathPart: "{serverId}"
  UsersResource:
    Type: AWS::ApiGateway::Resource
    Properties:
      RestApiId:
        Ref: CustomIdentityProviderApi
      ParentId:
        Ref: ServerIdResource
      PathPart: users
  UserNameResource:
    Type: AWS::ApiGateway::Resource
    Properties:
      RestApiId:
        Ref: CustomIdentityProviderApi
      ParentId:
        Ref: UsersResource
      PathPart: "{username}"
  GetUserConfigResource:
    Type: AWS::ApiGateway::Resource
    Properties:
      RestApiId:
        Ref: CustomIdentityProviderApi
      ParentId:
        Ref: UserNameResource
      PathPart: config
  GetUserConfigRequest:
    Type: AWS::ApiGateway::Method
    DependsOn: GetUserConfigResponseModel
    Properties:
      AuthorizationType: AWS_IAM
      HttpMethod: GET
      Integration:
        Type: AWS
        IntegrationHttpMethod: POST
        Uri:
          Fn::Join:
            - ""
            - - "arn:"
              - Ref: AWS::Partition
              - ":apigateway:"
              - Ref: AWS::Region
              - ":lambda:path/2015-03-31/functions/"
              - Fn::GetAtt: UserAuthenticationLambda.Arn
              - "/invocations"
        IntegrationResponses:
        - StatusCode: "200"
        RequestTemplates:
          application/json: |
            {
              "username": "$util.urlDecode($input.params('username'))",
              "password": "$util.escapeJavaScript($util.base64Decode($input.params('PasswordBase64'))).replaceAll("\\'","'")",
              "protocol": "$input.params('protocol')",
              "serverId": "$input.params('serverId')",
              "sourceIp": "$input.params('sourceIp')"
            }
      RequestParameters:
        method.request.header.PasswordBase64: false
        method.request.querystring.protocol: false
        method.request.querystring.sourceIp: false
      ResourceId:
        Ref: GetUserConfigResource
      RestApiId:
        Ref: CustomIdentityProviderApi
      MethodResponses:
      - StatusCode: "200"
        ResponseModels:
          application/json: UserConfigResponseModel
  GetUserConfigResponseModel:
    Type: AWS::ApiGateway::Model
    Properties:
      RestApiId:
        Ref: CustomIdentityProviderApi
      ContentType: application/json
      Description: API response for GetUserConfig
      Name: UserConfigResponseModel
      Schema:
        "$schema": http://json-schema.org/draft-04/schema#
        title: UserUserConfig
        type: object
        properties:
          HomeDirectory:
            type: string
          Role:
            type: string
          Policy:
            type: string
          PublicKeys:
            type: array
            items:
              type: string
  
  

  TransferS3AccessPolicy:
    Type: 'AWS::IAM::Policy'
    Properties:
      PolicyName: 
        Fn::Sub: ${Prefix}-TransferS3AccessPolicy
      Roles: 
        - Ref: TransferS3AccessRole
      PolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Action:
              - s3:ListBucket
            Resource: 
              - Fn::GetAtt: TransferS3Bucket.Arn
              - Fn::GetAtt: TransferPublicKeyS3Bucket.Arn
          - Effect: Allow
            Action:
              - s3:PutObject
              - s3:PutObjectAcl
              - s3:GetObject
              - s3:GetObjectAcl
              - s3:GetObjectVersion
              - s3:DeleteObject
              - s3:DeleteObjectVersion
              - s3:GetBucketLocation
            Resource: 
              - Fn::GetAtt: TransferS3Bucket.Arn
              - Fn::Sub:
                  - ${BucketARN}/*
                  - BucketARN:
                      Fn::GetAtt: TransferS3Bucket.Arn
              - Fn::GetAtt: TransferPublicKeyS3Bucket.Arn
              - Fn::Sub:
                  - ${KeysBucketARN}/*
                  - KeysBucketARN:
                      Fn::GetAtt: TransferPublicKeyS3Bucket.Arn
  
  UserAuthenticationLambdaExecutionPolicy:
    Type: 'AWS::IAM::Policy'
    Properties:
      PolicyName: 
        Fn::Sub: ${Prefix}-UserAuthentication-LambdaExecutionPolicy
      Roles:
        - Ref: UserAuthenticationLambdaExecutionRole
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Action:
              - cognito-idp:InitiateAuth
            Resource:
              - Fn::Join:
                  - ""
                  - - Fn::Sub: "arn:aws:cognito-idp:${AWS::Region}:${AWS::AccountId}:userpool/"
                    - Ref: UserPool
          - Effect: Allow
            Action: 
              - ssm:GetParametersByPath
              - ssm:GetParameter
              - ssm:GetParameters
            Resource:
              - Fn::Sub: arn:aws:ssm:${AWS::Region}:${AWS::AccountId}:parameter${TransferSSMParameter}/${Prefix}*
          - Effect: Allow
            Action: 
              - s3:GetObject
              - s3:ListBucket
            Resource:
              - Fn::GetAtt: TransferPublicKeyS3Bucket.Arn
              - Fn::Sub:
                  - ${BucketARN}/*
                  - BucketARN:
                      Fn::GetAtt: TransferPublicKeyS3Bucket.Arn
  
  
  
  UserAuthenticationLambdaPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:invokeFunction
      FunctionName:
        Ref: UserAuthenticationLambda
      Principal: apigateway.amazonaws.com
      SourceArn:
        Fn::Join:
          - ''
          - - Fn::Sub: "arn:${AWS::Partition}:execute-api:${AWS::Region}:${AWS::AccountId}:"
            - Ref: CustomIdentityProviderApi
            - "/*"
  
  TransferCanReadApiPolicy:
    Type: 'AWS::IAM::Policy'
    Properties:
      PolicyName: 
        Fn::Sub: ${Prefix}-TransferCanReadApiPolicy
      Roles:
        - Ref: TransferIdentityProviderRole
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Action:
            - apigateway:GET
            Resource: 
              Fn::Join:
                - ''
                - - Fn::Sub: "arn:${AWS::Partition}:execute-api:${AWS::Region}:${AWS::AccountId}:"
                  - Ref: CustomIdentityProviderApi
                  - Fn::Sub: "/${Prefix}/GET/servers/*/users/*"
  TransferCanInvokeApiPolicy:
    Type: 'AWS::IAM::Policy'
    Properties:
      PolicyName: 
        Fn::Sub: ${Prefix}-TransferCanInvokeApiPolicy
      Roles:
        - Ref: TransferIdentityProviderRole
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Action:
            - execute-api:Invoke
            Resource:
              Fn::Join:
                - ''
                - - Fn::Sub: "arn:${AWS::Partition}:execute-api:${AWS::Region}:${AWS::AccountId}:"
                  - Ref: CustomIdentityProviderApi
                  - Fn::Sub: "/${Prefix}/GET/servers/*/users/*"

Outputs:
  TransferS3BucketName:
    Description: S3 Bucket Access for Transfer Service Bucket Name
    Value: 
      Ref: TransferS3Bucket
    Export:
      Name: 
        Fn::Sub: ${Prefix}-TransferS3BucketName
  TransferPublicKeyS3BucketName:
    Description: S3 Bucket Access for Transfer Public Key Bucket ARN
    Value: 
      Ref: TransferPublicKeyS3Bucket
    Export:
      Name: 
        Fn::Sub: ${Prefix}-TransferPublicKeyS3BucketName
  UserPool:
    Description: FTP Cognito User Pool
    Value: 
      Ref: UserPool
    Export:
      Name: 
        Fn::Sub: ${Prefix}-UserPool
  UserPoolClient:
    Description: FTP Cognito User Pool App Client
    Value: 
      Ref: UserPoolClient
    Export:
      Name: 
        Fn::Sub: ${Prefix}-UserPoolClient
  TransferServerId:
    Description: Transfer Server
    Value: 
      Fn::GetAtt: TransferServer.ServerId
    Export:
      Name: 
        Fn::Sub: ${Prefix}-TransferServer
