---
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
  TransferIdentityProviderUrl:
    Description: URL to pass to AWS Transfer CreateServer call as part of optional IdentityProviderDetails
    Value:
      Fn::Join:
        - ''
        - - 'https://'
          - Ref: OktaIdentityProviderApi
          - '.execute-api.'
          - Ref: AWS::Region
          - '.'
          - Ref: AWS::URLSuffix
          - '/'
          - Ref: ApiStage
    Condition: NotCreateServer
  TransferIdentityProviderInvocationRole:
    Description: IAM Role to pass to AWS Transfer CreateServer call as part of optional IdentityProviderDetails
    Value:
      Fn::GetAtt: TransferIdentityProviderRole.Arn
    Condition: NotCreateServer
Resources:
  TransferServer:
    Type: AWS::Transfer::Server
    Condition: CreateServer
    Properties:
      EndpointType: PUBLIC
      IdentityProviderDetails:
        InvocationRole:
          Fn::GetAtt: TransferIdentityProviderRole.Arn
        Url:
          Fn::Join:
            - ''
            - - 'https://'
              - Ref: OktaIdentityProviderApi
              - '.execute-api.'
              - Ref: AWS::Region
              - '.'
              - Ref: AWS::URLSuffix
              - '/'
              - Ref: ApiStage
      IdentityProviderType: API_GATEWAY
      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}
  OktaIdentityProviderApi:
    Type: AWS::ApiGateway::RestApi
    Properties:
      Name: Transfer Okta Identity Provider basic template API
      Description: API used for GetUserConfig requests
      FailOnWarnings: true
      EndpointConfiguration:
        Types:
        - REGIONAL
  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
  ApiCloudWatchLogsRole:
    Description: IAM role used by API Gateway to log API requests to CloudWatch
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: apigateway.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: ApiGatewayLogsPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:DescribeLogGroups
                  - logs:DescribeLogStreams
                  - logs:DescribeQueries
                  - logs:FilterLogEvents
                  - logs:GetLogEvents
                  - logs:GetLogGroupFields
                  - logs:GetLogRecord
                  - logs:GetQueryResults
                  - logs:PutLogEvents
                  - logs:StartQuery
                  - logs:StopQuery
                Resource: '*'
  ApiLoggingAccount:
    Type: AWS::ApiGateway::Account
    DependsOn: OktaIdentityProviderApi
    Properties:
      CloudWatchRoleArn:
        Fn::GetAtt: ApiCloudWatchLogsRole.Arn
  ApiStage:
    Type: AWS::ApiGateway::Stage
    Properties:
      DeploymentId:
        Ref: ApiDeployment202008
      MethodSettings:
        - DataTraceEnabled: false
          HttpMethod: '*'
          LoggingLevel: INFO
          ResourcePath: '/*'
      RestApiId:
        Ref: OktaIdentityProviderApi
      StageName: prod
  ApiDeployment202008:
    DependsOn: GetUserConfigRequest
    Type: AWS::ApiGateway::Deployment
    Properties:
      RestApiId:
        Ref: OktaIdentityProviderApi
  TransferIdentityProviderRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: transfer.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: TransferCanInvokeThisApi
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action: execute-api:Invoke
                Resource:
                  Fn::Sub: arn:${AWS::Partition}:execute-api:${AWS::Region}:${AWS::AccountId}:${OktaIdentityProviderApi}/prod/GET/*
        - PolicyName: TransferCanReadThisApi
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action: apigateway:GET
                Resource: '*'
  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: apigateway.amazonaws.com
      SourceArn:
        Fn::Join:
          - ''
          - - 'arn:aws:execute-api:'
            - Ref: AWS::Region
            - ':'
            - Ref: AWS::AccountId
            - ':'
            - Ref: OktaIdentityProviderApi
            - '/*'
  ServersResource:
    Type: AWS::ApiGateway::Resource
    Properties:
      RestApiId:
        Ref: OktaIdentityProviderApi
      ParentId:
        Fn::GetAtt:
          - OktaIdentityProviderApi
          - RootResourceId
      PathPart: servers
  ServerIdResource:
    Type: AWS::ApiGateway::Resource
    Properties:
      RestApiId:
        Ref: OktaIdentityProviderApi
      ParentId:
        Ref: ServersResource
      PathPart: '{serverId}'
  UsersResource:
    Type: AWS::ApiGateway::Resource
    Properties:
      RestApiId:
        Ref: OktaIdentityProviderApi
      ParentId:
        Ref: ServerIdResource
      PathPart: users
  UserNameResource:
    Type: AWS::ApiGateway::Resource
    Properties:
      RestApiId:
        Ref: OktaIdentityProviderApi
      ParentId:
        Ref: UsersResource
      PathPart: '{username}'
  GetUserConfigResource:
    Type: AWS::ApiGateway::Resource
    Properties:
      RestApiId:
        Ref: OktaIdentityProviderApi
      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:aws:apigateway:'
              - Ref: AWS::Region
              - ':lambda:path/2015-03-31/functions/'
              - Fn::GetAtt:
                - GetUserConfigLambda
                - 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: OktaIdentityProviderApi
      MethodResponses:
        - StatusCode: 200
          ResponseModels:
            application/json: UserConfigResponseModel
  GetUserConfigResponseModel:
    Type: AWS::ApiGateway::Model
    Properties:
      RestApiId:
        Ref: OktaIdentityProviderApi
      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
