---
AWSTemplateFormatVersion: '2010-09-09'
Description: A basic template for creating a Lambda-backed API Gateway for integrating with OKTA as
  a custom identity provider along with MultiFactor Authentication in AWS Transfer.
Parameters:
  AWSTransferS3Bucket:
    Type: String
    Description: The name of the S3 bucket used for the server
  OKTAAuthAPIURL:
    Type: String
    Description: Okta Authorization API URL (such as https://<YOUR-OKTA-ENDPOINT>/api/v1/authn)
  UserSigninDomain:
    Type: String
    Description: OKTA User Sign-in domain (for example, for joesmith@amazon.com, the domain will be amazon.com)
  CreateServer:
    AllowedValues:
      - 'true'
      - 'false'
    Type: String
    Description: Whether this stack creates a server internally or not. If a server is created internally,
      the customer identity provider is automatically associated with it.
    Default: 'true'
Conditions:
  CreateServer:
    Fn::Equals:
      - Ref: CreateServer
      - 'true'
  NotCreateServer:
    Fn::Not:
      - Condition: CreateServer
Outputs:
  ServerId:
    Value:
      Fn::GetAtt: TransferServer.ServerId
    Condition: CreateServer
  StackArn:
    Value:
      Ref: AWS::StackId
Resources:
  TransferServer:
    Type: AWS::Transfer::Server
    Condition: CreateServer
    Properties:
      EndpointType: PUBLIC
      IdentityProviderDetails:
        Function:
          Fn::GetAtt: GetUserConfigLambda.Arn
      IdentityProviderType: AWS_LAMBDA
      LoggingRole:
        Fn::GetAtt: CloudWatchLoggingRole.Arn
  CloudWatchLoggingRole:
    Description: IAM role used by AWS Transfer to log user activity to CloudWatch
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: transfer.amazonaws.com
            Action: sts:AssumeRole
      Path: '/'
      Policies:
        - PolicyName: TransferLogsPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              Effect: Allow
              Action:
                - logs:CreateLogGroup
                - logs:CreateLogStream
                - logs:DescribeLogStreams
                - logs:PutLogEvents
              Resource: '*'
  TransferS3Access: # Allows AWS Transfer to access S3 bucket
    Description: Allows AWS Transfer access S3 bucket
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: transfer.amazonaws.com
            Action: sts:AssumeRole
      Path: '/'
      Policies:
        - PolicyName: TransferS3BucketAccessPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - s3:DeleteObject
                  - s3:PutObject
                  - s3:GetObject
                  - s3:DeleteObjectVersion
                  - s3:GetObjectVersion
                  - s3:GetObjectACL
                  - s3:PutObjectACL
                Resource:
                  Fn::Sub: arn:aws:s3:::${AWSTransferS3Bucket}/*
              - Effect: Allow
                Action:
                  - s3:ListBucket
                  - s3:GetBucketLocation
                Resource:
                  Fn::Sub: arn:aws:s3:::${AWSTransferS3Bucket}
  LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
      - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
  GetUserConfigLambda:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile:
          Fn::Sub: |
            import boto3
            import json
            import logging
            import os
            from urllib import request, parse

            logger = logging.getLogger()
            logger.setLevel(logging.INFO)

            def lambda_handler(event, context):
                logger.info("Event: User Name [{}], Server ID: [{}]".format(event["username"], event["serverId"]))

                response = {}

                if "username" not in event or "serverId" not in event or "password" not in event:
                    return response

                # It is recommended to verify server ID against some value, this template does not verify server ID
                server_id = event["serverId"]
                username = event["username"]
                password = event["password"]

                # This code is designed for usage with Google Authenticator. Google Authenticator has a 6 digit token.
                mfaToken = password[-6:]
                password = password[:-6]

                status_code = auth_with_okta(get_full_username(username), password, mfaToken)

                if status_code == 200:
                    home_directory = "${AWSTransferS3Bucket}" + "/" + username + "/"
                    response["Role"] = "${TransferS3Access.Arn}"
                    response["HomeDirectory"] = "/" + home_directory

                    # Optional JSON blob to further restrict this user's permissions
                    response["Policy"] = ""
                    logger.info("Message: {}".format(response))
                else:
                    logger.info("Failed to authenticate user [{}] with Okta. Received status code of {}".format(get_full_username(username), status_code))

                return response

            def get_full_username(user_name):
                return user_name + "@" + "${UserSigninDomain}"

            # This function authenticates a user with Okta and returns a status code
            def auth_with_okta(user_name, password, mfaToken):

                headers = {"Accept": "application/json", "Content-Type": "application/json"}
                primary_credentials = {"username": user_name, "password": password}

                try:
                    primary_request = request.Request(url="${OKTAAuthAPIURL}", data=json.dumps(primary_credentials).encode('utf-8'), headers=headers)
                    primary_response = request.urlopen(primary_request)

                    logger.info("Okta response: [{}]".format(primary_response.status))
                except Exception as e:
                    logger.info("Error authing with Okta: {}".format(str(e)))
                    return 0

                primary = json.loads(primary_response.read().decode('utf-8'))
                stateToken = primary['stateToken']
                secondary_url = primary['_embedded']['factors'][0]['_links']['verify']['href']
                secondary_credentials = {"stateToken": stateToken, "passCode": mfaToken}

                try:
                    secondary_resquest = request.Request(url=secondary_url, data=json.dumps(secondary_credentials).encode('utf-8'), headers=headers)
                    secondary_response = request.urlopen(secondary_resquest)

                    logger.info("Okta response: [{}]".format(secondary_response.status))
                    return secondary_response.status
                except Exception as e:
                    logger.info("Error authing with Okta: {}".format(str(e)))
                    return 0
      Description: A function to provide IAM roles and policies for given user and serverId.
      Handler: index.lambda_handler
      Role:
        Fn::GetAtt: LambdaExecutionRole.Arn
      Runtime: python3.11
      Timeout: 15
  GetUserConfigLambdaPermission:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:invokeFunction
      FunctionName:
        Fn::GetAtt: GetUserConfigLambda.Arn
      Principal: transfer.amazonaws.com
      SourceArn:
        Fn::If:
          - CreateServer
          - Fn::GetAtt: TransferServer.Arn
          - Fn::Sub: arn:${AWS::Partition}:transfer:${AWS::Region}:${AWS::AccountId}:server/*
