---
AWSTemplateFormatVersion: '2010-09-09'
Description: CloudFormation that uses AWS Lambda with an AWS Transfer Family server
  to integrate SecretsManager as an identity provider. It authenticates against an
  entry in AWS Secrets Manager of the format aws/transfer/username. The secrets manager also
  stores MFA Authentication code for the SFTP users
Metadata:
  AWS::CloudFormation::Interface:
    ParameterGroups:
      -
        Label:
          default: "AWS Transfer Family for SFTP Configuration"
        Parameters:
          - ExistingS3Bucket
          - UserName
          - Password
          - MFAAuthCode
          - PublicOrPrivate
      -
        Label:
          default: "VPC and Network Configuration for Private Deployment"
        Parameters:
          - VPCID
          - SubnetID
          - SecurityGroupID
    ParameterLabels:
      PublicOrPrivate:
        default: "Public or private"
      VPCID:
        default: "Enter VPC ID of AWS Transfer Family Server deployment"
      SubnetID:
        default: "Enter list of Subnet IDs (separated by comma) where AWS Transfer Family Server VPC endpoint will be deployed to"
      SecurityGroupID:
        default: "Enter list of Security Group IDs (separated by comma) for the VPC endpoint"
Parameters:
  ExistingS3Bucket:
    Type: String
    Description: Existing S3 Bucket name that will be used as storage for AWS Transfer Family for SFTP. If left blank, a new S3 bucket will be created.
    Default: ''
  UserName:
    Type: String
    Description: User name that will be created for AWS Transfer Family for SFTP
    Default: 'demo-user01'
  Password:
    Type: String
    NoEcho: true
    Description: User password that will be created for AWS Transfer Family for SFTP
  MFAAuthCode:
    Type: String
    NoEcho: true
    AllowedPattern: '[A-Z2-7]{16,}'
    Description: User MFA Secret code that is registered in MFA device. It must be a string with a minimum length of 16 characters consisting only alphabet A-Z and 2-7 as specified in https://datatracker.ietf.org/doc/html/rfc4648#section-6
  PublicOrPrivate:
    Type: String
    Default: 'Public'
    AllowedValues:
      - Public
      - Private
    Description: Specify whether the AWS Transfer Family will be accessible publicly from Internet or only privately accessible within VPC. If Public is selected, please ignore VPC and network section below.
  VPCID:
    Type: String
    Default: ''
  SubnetID:
    Type: CommaDelimitedList
    Default: ''
  SecurityGroupID:
    Type: CommaDelimitedList
    Default: ''

Conditions:
  IsPublic: !Equals [!Ref PublicOrPrivate, 'Public']
  IsNewS3Bucket: !Equals [!Ref ExistingS3Bucket, '']
  IsExistingS3Bucket: !Not [!Equals [!Ref ExistingS3Bucket, '']]

Resources:
# Create home S3 bucket
  HomeS3Bucket:
    Condition: IsNewS3Bucket
    Type: AWS::S3::Bucket
    Properties:
      BucketName: !Join ["-", ["transfer-home", !Select [0, !Split ["-", !Select [2, !Split ["/", !Ref "AWS::StackId"]]]]]]
      PublicAccessBlockConfiguration:
        BlockPublicAcls: true
        BlockPublicPolicy: true
        IgnorePublicAcls: true
        RestrictPublicBuckets: true
      Tags:
        - Key: "blogSource"
          Value: "TransferFamilyMFA"
# Create a secret in Secrets Manager
  TransferSecret:
    Type: AWS::SecretsManager::Secret
    Properties:
      Name: !Join ["/", [!Join ["-", ["transfer-user", !Select [0, !Split ["-", !Select [2, !Split ["/", !Ref "AWS::StackId"]]]]]], !Ref UserName]]
      SecretString:
        Fn::Sub: '{"username":"${UserName}","password":"${Password}","mfaAuthCode":"${MFAAuthCode}"}'
      Tags:
        - Key: "blogSource"
          Value: "TransferFamilyMFA"
# Upload user home folder object to home S3 bucket
  S3PutObjectFunctionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Principal:
            Service:
            - lambda.amazonaws.com
          Action:
          - sts:AssumeRole
      Path: "/"
      ManagedPolicyArns:
      - Fn::Sub: arn:${AWS::Partition}:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        'Fn::If':
          - IsNewS3Bucket
          - - PolicyName: S3Writer
              PolicyDocument:
                Version: '2012-10-17'
                Statement:
                - Effect: Allow
                  Action:
                  - s3:PutObject
                  - s3:GetObject
                  - s3:ListBucket
                  Resource:
                    - !Sub arn:aws:s3:::${HomeS3Bucket}/*
                    - !Sub arn:aws:s3:::${HomeS3Bucket}
          - - PolicyName: S3Writer
              PolicyDocument:
                Version: '2012-10-17'
                Statement:
                - Effect: Allow
                  Action:
                  - s3:PutObject
                  - s3:GetObject
                  - s3:ListBucket
                  Resource:
                    - !Sub arn:aws:s3:::${ExistingS3Bucket}/*
                    - !Sub arn:aws:s3:::${ExistingS3Bucket}
      Tags:
      - Key: "blogSource"
        Value: "TransferFamilyMFA"
  S3PutObjectFunction:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile:
          Fn::Sub: |
            import boto3
            import cfnresponse
            def handler(event, context):
              try:
                response_data = {}
                if (event['RequestType'] == 'Create'):
                  the_bucket = event['ResourceProperties']['Bucket']
                  dir_name = event['ResourceProperties']['Key']
                  print("Event received ", str(event))
                  print("bucket received ", the_bucket)
                  print("key received ", dir_name)
                  s_3 = boto3.client('s3')
                  s_3.put_object(Bucket=the_bucket,Key=dir_name)
                cfnresponse.send(event,context,cfnresponse.SUCCESS,response_data)
              except Exception as e:
                response_data['Data'] = str(e)
                cfnresponse.send(event,context,cfnresponse.FAILED,response_data)
      Description: 'Used to put objects into S3.'
      Handler: index.handler
      Role:
        Fn::GetAtt: S3PutObjectFunctionRole.Arn
      Runtime: python3.11
      Timeout: 30
      ReservedConcurrentExecutions: 5
      Tags:
        - Key: "blogSource"
          Value: "TransferFamilyMFA"
    DependsOn: S3PutObjectFunctionRole
  CreateTransferUserHomeFolder:
    Type: Custom::S3PutObject
    Properties:
      ServiceToken:
        Fn::GetAtt: S3PutObjectFunction.Arn
      Bucket:
        'Fn::If':
          - IsNewS3Bucket
          - !Ref HomeS3Bucket
          - !Ref ExistingS3Bucket
      Key:
        Fn::Sub: ${UserName}/
      Tags:
        - Key: "blogSource"
          Value: "TransferFamilyMFA"
  CloudWatchTransferLoggingRole:
    Description: IAM role used by Transfer to log API requests to CloudWatch
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - transfer.amazonaws.com
            Action:
              - sts:AssumeRole
      Policies:
        - PolicyName: CloudWatchTransferLogsPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:DescribeLogStreams
                  - logs:PutLogEvents
                Resource:
                  Fn::Sub: '*'
      Tags:
        - Key: "blogSource"
          Value: "TransferFamilyMFA"
  LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      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/AWSLambdaBasicExecutionRole
      Policies:
      - PolicyName: LambdaSecretsPolicy
        PolicyDocument:
          Version: '2012-10-17'
          Statement:
          - Effect: Allow
            Action:
            - secretsmanager:GetSecretValue
            Resource: !Sub
              - arn:${AWS::Partition}:secretsmanager:${AWS::Region}:${AWS::AccountId}:secret:${SecretsManagerPrefix}/*
              - SecretsManagerPrefix: !Join ["-", ["transfer-user", !Select [0, !Split ["-", !Select [2, !Split ["/", !Ref "AWS::StackId"]]]]]]
      Tags:
        - Key: "blogSource"
          Value: "TransferFamilyMFA"
  TransferAccessS3Role:
    Description: IAM role used by Transfer to Access S3 and contain trust relationship while servicing user request
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service:
              - transfer.amazonaws.com
            Action:
              - sts:AssumeRole
      Policies:
        'Fn::If':
          - IsNewS3Bucket
          - - PolicyName: TransferAccessS3RolePolicy
              PolicyDocument:
                Version: '2012-10-17'
                Statement:
                - Effect: Allow
                  Action:
                  - s3:PutObject
                  - s3:PutObjectAcl
                  - s3:GetObject
                  - s3:GetObjectAcl
                  - s3:GetObjectVersion
                  - s3:DeleteObject
                  - s3:DeleteObjectVersion
                  - s3:GetBucketLocation
                  - s3:ListBucket
                  Resource:
                    - !Sub arn:aws:s3:::${HomeS3Bucket}/*
                    - !Sub arn:aws:s3:::${HomeS3Bucket}
          - - PolicyName: TransferAccessS3RolePolicy
              PolicyDocument:
                Version: '2012-10-17'
                Statement:
                - Effect: Allow
                  Action:
                  - s3:PutObject
                  - s3:PutObjectAcl
                  - s3:GetObject
                  - s3:GetObjectAcl
                  - s3:GetObjectVersion
                  - s3:DeleteObject
                  - s3:DeleteObjectVersion
                  - s3:GetBucketLocation
                  - s3:ListBucket
                  Resource:
                    - !Sub arn:aws:s3:::${ExistingS3Bucket}/*
                    - !Sub arn:aws:s3:::${ExistingS3Bucket}
      Tags:
        - Key: "blogSource"
          Value: "TransferFamilyMFA"
  TransferServer:
    Type: AWS::Transfer::Server
    Properties:
      EndpointType:
        'Fn::If':
          - IsPublic
          - PUBLIC
          - VPC
      EndpointDetails:
        'Fn::If':
          - IsPublic
          - !Ref 'AWS::NoValue'
          - VpcId:
              'Fn::If':
                - IsPublic
                - !Ref 'AWS::NoValue'
                - !Ref VPCID
            SubnetIds:
              'Fn::If':
                - IsPublic
                - !Ref 'AWS::NoValue'
                - !Ref SubnetID
            SecurityGroupIds:
              'Fn::If':
                - IsPublic
                - !Ref 'AWS::NoValue'
                - !Ref SecurityGroupID
      IdentityProviderDetails:
        Function:
          Fn::GetAtt: TransferFamilyAuthLambda.Arn
      IdentityProviderType: AWS_LAMBDA
      LoggingRole:
        Fn::GetAtt: CloudWatchTransferLoggingRole.Arn
      Tags:
        - Key: "blogSource"
          Value: "TransferFamilyMFA"
  TransferFamilyAuthLambda:
    Type: AWS::Lambda::Function
    Properties:
      Description: 'Lambda function that handles authentication with MFA'
      Environment:
        Variables:
          'Fn::If':
            - IsNewS3Bucket
            - BucketName : !Ref HomeS3Bucket
              SecretsManagerPrefix: !Join ["-", ["transfer-user", !Select [0, !Split ["-", !Select [2, !Split ["/", !Ref "AWS::StackId"]]]]]]
            - BucketName : !Ref ExistingS3Bucket
              SecretsManagerPrefix: !Join ["-", ["transfer-user", !Select [0, !Split ["-", !Select [2, !Split ["/", !Ref "AWS::StackId"]]]]]]
      Handler: index.lambda_handler
      Role:
        Fn::GetAtt: LambdaExecutionRole.Arn
      Runtime: python3.11
      Timeout: 30
      ReservedConcurrentExecutions: 100
      Code:
        ZipFile:
          Fn::Sub: |
            import boto3
            import json
            import os
            import base64
            import hmac, struct, hashlib, time
            from botocore.exceptions import ClientError

            BUCKET_NAME = os.environ["BucketName"]
            SECRETS_MANAGER_PREFIX = os.environ["SecretsManagerPrefix"]

            def lambda_handler(event, context):
              auth_response = {}
              data_ret = {}
              username = event["username"]
              passwordAndMFA = event["password"]
              secret = getSecret(username)
              strLength = len(passwordAndMFA)
              password = passwordAndMFA[0:strLength-6]
              mfaInput = passwordAndMFA[-6:]

              mfaSecret = get_totp_token(secret["mfaAuthCode"])

              if (secret["username"] == username and secret["password"] == password and mfaSecret == mfaInput):
                directorymapping = [{"Entry": "/", "Target": "/"+ BUCKET_NAME + "/" + username}]

                data_ret["Role"] = "${TransferAccessS3Role.Arn}"
                data_ret["HomeDirectoryType"] = "LOGICAL"
                data_ret["HomeDirectoryDetails"] = json.dumps(directorymapping)
              return data_ret

            def getSecret(username):
              secret_name = SECRETS_MANAGER_PREFIX + "/" + username
              region_name = "${AWS::Region}"

              # Create a Secrets Manager client
              session = boto3.session.Session()
              client = session.client(
                  service_name='secretsmanager',
                  region_name=region_name
              )
              try:
                  get_secret_value_response = client.get_secret_value(
                      SecretId=secret_name
                  )
              except ClientError as e:
                  print("secret error")
                  print(e)
                  if e.response['Error']['Code'] == 'DecryptionFailureException':
                      # Secrets Manager can't decrypt the protected secret text using the provided KMS key.
                      # Deal with the exception here, and/or rethrow at your discretion.
                      raise e
                  elif e.response['Error']['Code'] == 'InternalServiceErrorException':
                      # An error occurred on the server side.
                      # Deal with the exception here, and/or rethrow at your discretion.
                      raise e
                  elif e.response['Error']['Code'] == 'InvalidParameterException':
                      # You provided an invalid value for a parameter.
                      # Deal with the exception here, and/or rethrow at your discretion.
                      raise e
                  elif e.response['Error']['Code'] == 'InvalidRequestException':
                      # You provided a parameter value that is not valid for the current state of the resource.
                      # Deal with the exception here, and/or rethrow at your discretion.
                      raise e
                  elif e.response['Error']['Code'] == 'ResourceNotFoundException':
                      # We can't find the resource that you asked for.
                      # Deal with the exception here, and/or rethrow at your discretion.
                      raise e
              else:
                  # Decrypts secret using the associated KMS key.
                  # Depending on whether the secret is a string or binary, one of these fields will be populated.

                  if 'SecretString' in get_secret_value_response:
                      secret = get_secret_value_response['SecretString']
                      return json.loads(secret)
                  else:
                      decoded_binary_secret = base64.b64decode(get_secret_value_response['SecretBinary'])
                      return decoded_binary_secret

            # Following Google Authenticator routine codes are originated from https://github.com/grahammitchell/google-authenticator
            def get_hotp_token(secret, intervals_no):
              key = base64.b32decode(normalize(secret), True)
              msg = struct.pack(">Q", intervals_no)
              h = bytearray(hmac.new(key, msg, hashlib.sha1).digest())
              o = h[19] & 15
              h = str((struct.unpack(">I", h[o:o+4])[0] & 0x7fffffff) % 1000000)
              return prefix0(h)

            def get_totp_token(secret):
              return get_hotp_token(secret, intervals_no=int(time.time())//30)

            def normalize(key):
              k2 = key.strip().replace(' ','')
              if len(k2)%8 != 0:
                k2 += '='*(8-len(k2)%8)
              return k2

            def prefix0(h):
              if len(h) < 6:
                h = '0'*(6-len(h)) + h
              return h
      Tags:
        - Key: "blogSource"
          Value: "TransferFamilyMFA"
  TransferFamilyAuthLambdaResourceBasedPolicy:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !GetAtt TransferFamilyAuthLambda.Arn
      Action: lambda:InvokeFunction
      Principal: transfer.amazonaws.com
      SourceArn: !GetAtt TransferServer.Arn
Outputs:
  ServerId:
    Description: "Transfer Family Server ID"
    Value:
      Fn::GetAtt: TransferServer.ServerId
  TransferFamilyServer:
    Description: "Transfer Family Server Endpoint"
    Value:
      Fn::Sub: ${TransferServer.ServerId}.server.transfer.${AWS::Region}.${AWS::URLSuffix}
  StackArn:
    Value:
      Ref: AWS::StackId
  S3Bucket:
    Description: "S3 Bucket Name"
    Value:
      'Fn::If':
        - IsNewS3Bucket
        - !Ref HomeS3Bucket
        - !Ref ExistingS3Bucket
  SecretsManagerPrefixName:
    Description: "Secrets Manager Prefix Name"
    Value: !Join ["/", [!Join ["-", ["transfer-user", !Select [0, !Split ["-", !Select [2, !Split ["/", !Ref "AWS::StackId"]]]]]], ""]]
