---
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:
          - ExistingEFSARN
          - UserName
          - Password
          - UID
          - GID
          - MFAAuthCode
          - PublicOrPrivate
      -
        Label:
          default: "VPC and Network Configuration for Private Deployment"
        Parameters:
          - VPCID
          - SubnetID
          - SecurityGroupID
    ParameterLabels:
      ExistingEFSARN:
        default: "Existing EFS ARN"
      UID:
        default: "POSIX User ID"
      GID:
        default: "POSIX Group ID"
      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:
  ExistingEFSARN:
    Type: String
    Description: ARN of Existing EFS that will be used as storage for AWS Transfer Family for SFTP. Example arn:aws:elasticfilesystem:<region>:<account-id>:file-system/<file-system-id>
    Default: ''
  UserName:
    Type: String
    Description: User name that will be created for AWS Transfer Family for SFTP. A folder with the same name (e.g. /demo-user01) with corresponding POSIX profile must have been created in the EFS access points
    Default: 'demo-user01'
  Password:
    Type: String
    NoEcho: true
    Description: User password that will be created for AWS Transfer Family for SFTP
  UID:
    Type: String
    Description: POSIX Profile User ID of the user in EFS. For example '1000'
    AllowedPattern: '[0-9]{1,10}'
  GID:
    Type: String
    Description: POSIX Profile Group ID of the user in EFS. For example '1000'
    AllowedPattern: '[0-9]{1,10}'
  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. For example 'ABCDEFGHIJ234567'
  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']

Resources:
# 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}","UID":"${UID}","GID":"${GID}"}'
      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"
  TransferAccessEFSRole:
    Description: IAM role used by Transfer to Access EFS 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:
          - PolicyName: TransferAccessEFSRolePolicy
            PolicyDocument:
              Version: '2012-10-17'
              Statement:
              - Effect: Allow
                Action:
                - elasticfilesystem:ClientMount
                - elasticfilesystem:ClientWrite
                Resource: !Ref ExistingEFSARN
      Tags:
        - Key: "blogSource"
          Value: "TransferFamilyMFA"
  TransferServer:
    Type: AWS::Transfer::Server
    Properties:
      Domain: 'EFS'
      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:
          EFSName :  !Select [1, !Split ["/", !Ref ExistingEFSARN]]
          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

            EFS_NAME = os.environ["EFSName"]
            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"])
              userUid = secret["UID"]
              userGid = secret["GID"]

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

                posix = "{\"Gid\": " + userGid + ",\"Uid\": " + userUid + ", \"SecondaryGids\": []}"

                data_ret["Role"] = "${TransferAccessEFSRole.Arn}"
                data_ret["PosixProfile"] = json.loads(posix)
                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
  SecretsManagerPrefixName:
    Description: "Secrets Manager Prefix Name"
    Value: !Join ["/", [!Join ["-", ["transfer-user", !Select [0, !Split ["-", !Select [2, !Split ["/", !Ref "AWS::StackId"]]]]]], ""]]
