AWSTemplateFormatVersion: 2010-09-09
Description: CloudFormation template to deploy Transfer for SFTP with a custom authorizor.
Resources:
# Data S3 Bucket
  DataBucket:
    Type: 'AWS::S3::Bucket'
    Properties:
      AccessControl: Private
      PublicAccessBlockConfiguration:
        BlockPublicAcls: true
        IgnorePublicAcls: true
        BlockPublicPolicy: true
        RestrictPublicBuckets: true
      BucketEncryption:
        ServerSideEncryptionConfiguration:
        - ServerSideEncryptionByDefault:
            SSEAlgorithm: AES256
      VersioningConfiguration:
        Status: Enabled
  TransferUserRole:
    Type: 'AWS::IAM::Role'
    Properties:
      Path: /
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - transfer.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Policies:
        - PolicyName: root
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action: 's3:*'
                Resource: '*'
  SampleData:
    Type: 'Custom::SampleData'
    Properties:
      ServiceToken: !GetAtt SampleDataLambda.Arn
    DependsOn:
      - DataBucket
      - SampleDataLambda
  SampleDataLambdaRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Statement:
          - Action:
              - 'sts:AssumeRole'
            Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
        Version: 2012-10-17
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
      Policies:
        - PolicyName: SampleDataBucketFullAccess
          PolicyDocument:
            Statement:
              - Action:
                  - 's3:*'
                Effect: Allow
                Resource:
                  - !Join [ '', [!GetAtt DataBucket.Arn, '*']]       
                  - !Join [ '', [!GetAtt DataBucket.Arn, '/*']]        
        - PolicyName: DynamoDBWriteAccess
          PolicyDocument:
            Statement:
              - Action:
                  - "dynamodb:BatchGetItem"
                  - "dynamodb:BatchWriteItem"
                  - "dynamodb:ConditionCheckItem"
                  - "dynamodb:PutItem"
                  - "dynamodb:DescribeTable"
                  - "dynamodb:DeleteItem"
                  - "dynamodb:GetItem"
                  - "dynamodb:Scan"
                  - "dynamodb:Query"
                  - "dynamodb:UpdateItem"
                Effect: Allow
                Resource:
                  - !GetAtt Users.Arn
                  - !GetAtt Products.Arn
                  - !GetAtt UserProducts.Arn
            Version: 2012-10-17 
  SampleDataLambda:
    Type: 'AWS::Lambda::Function'
    DependsOn:
      - DataBucket
      - Users
      - Products
      - UserProducts
    Properties:
      Code: 
          ZipFile: |
            from __future__ import print_function          
            import logging
            import os
            import boto3
            import json
            import time
            import random
            import string
            import urllib3

            SUCCESS = "SUCCESS"
            FAILED = "FAILED"

            http = urllib3.PoolManager()

            def send(event, context, responseStatus, responseData, physicalResourceId=None, noEcho=False, reason=None):
                responseUrl = event['ResponseURL']

                print(responseUrl)

                responseBody = {
                    'Status' : responseStatus,
                    'Reason' : reason or "See the details in CloudWatch Log Stream: {}".format(context.log_stream_name),
                    'PhysicalResourceId' : physicalResourceId or context.log_stream_name,
                    'StackId' : event['StackId'],
                    'RequestId' : event['RequestId'],
                    'LogicalResourceId' : event['LogicalResourceId'],
                    'NoEcho' : noEcho,
                    'Data' : responseData
                }

                json_responseBody = json.dumps(responseBody)

                print("Response body:")
                print(json_responseBody)

                headers = {
                    'content-type' : '',
                    'content-length' : str(len(json_responseBody))
                }

                try:
                    response = http.request('PUT', responseUrl, headers=headers, body=json_responseBody)
                    print("Status code:", response.status)

                except Exception as e:
                    print("send(..) failed executing http.request(..):", e)

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

            ACTION_CREATE = "Create"
            ACTION_UPDATE = "Update"
            ACTION_DELETE = "Delete"

            STATUS_SUCCESS = "SUCCESS"
            STATUS_FAILED = "FAILED"

            def _rand_string(l):
                return ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(l))

            def _gen_physical_resource_id(event):
                if "PhysicalResourceId" in event.keys():
                    logger.info("PhysicalResourceId present in event, Using that for response")
                    return event['PhysicalResourceId']
                else:
                    logger.info("No physical resource id returned, generating one...")
                    return event['StackId'].split('/')[1] + '_' + event['LogicalResourceId'] + '_' + _rand_string(8)

            def _process(event, context, action):
                subscriptions_bucket = os.environ['DataBucket']

                files = [
                    [subscriptions_bucket, 'john/product-abc.txt'],
                    [subscriptions_bucket, 'jane/product-xyz.txt'],
                    [subscriptions_bucket, 'indices/index1-2018.txt'],
                    [subscriptions_bucket, 'indices/index2-2018.txt'],
                    [subscriptions_bucket, 'equities/equity1-2018.txt'],
                    [subscriptions_bucket, 'credit/credit1-2019.txt'],
                    [subscriptions_bucket, 'equities/equity1-2019.txt'],
                    [subscriptions_bucket, 'equities/equity2-2019.txt'],
                    [subscriptions_bucket, 'indices/index1-2019.txt'],
                    [subscriptions_bucket, 'indices/index2-2019.txt'],
                    [subscriptions_bucket, 'indices/index3-2019.txt'],
                ]

                s3 = boto3.client('s3')

                if action == ACTION_CREATE:
                    _adduserdata()
                    file_text = "Test data generated by CloudFormation stack %s. These objects will be automatically deleted on stack cleanup." % event['StackId']
                    for bucket, key in files:
                        print ("Putting data to s3://{}/{}".format (bucket, key))
                        s3.put_object(Bucket=bucket, Key=key, Body=file_text)
                elif action == ACTION_DELETE:
                    s3_res = boto3.resource('s3')
                    for bucket, key in files:
                        bucket_obj = s3_res.Bucket(bucket)
                        bucket_obj.object_versions.all().delete()


                response_body = {
                    'Status': STATUS_SUCCESS,
                    'PhysicalResourceId': _gen_physical_resource_id(event),
                    'StackId': event['StackId'],
                    'RequestId': event['RequestId'],
                    'LogicalResourceId': event['LogicalResourceId'],
                    'Reason': "",
                    'Data': {},
                }

                send(event, context, response_body["Status"],response_body["Data"])

            def _adduserdata():
                usersTableName = os.environ['UsersTableName'].split('/')[1]
                productsTableName = os.environ['ProductsTableName'].split('/')[1]
                userProductsTableName = os.environ['UserProductsTableName'].split('/')[1]
                dynamodb = boto3.resource('dynamodb')
                users_table = dynamodb.Table(usersTableName)
                user_1 = {"user_id":1, "first_name":'jane', "last_name":'doe', "email":'jane.doe@xyz.com', "user_name":'jane'}
                user_2 = {"user_id":2, "first_name":'john', "last_name":'doe', "email":'john.doe@xyz.com', "user_name":'john'}
            
                users_to_add = [user_1, user_2]

                with users_table.batch_writer() as batch:
                    for item in users_to_add:
                        response = batch.put_item(Item=item)

                products_table = dynamodb.Table(productsTableName)
                product_1 = {"product_id":1, "product_name":'equities'}
                product_2 = {"product_id":2, "product_name":'credit'}
                product_3 = {"product_id":3, "product_name":'indices'}

                products_to_add = [product_1, product_2, product_3]

                with products_table.batch_writer() as batch:
                    for item in products_to_add:
                        response = batch.put_item(Item=item)   

                userproducts_table = dynamodb.Table(userProductsTableName)
                userproduct_1 = {"user_id":1, "product_id":1,"relation":"P", "user_name":"jane","product_name":"equities"}
                userproduct_2 = {"user_id":1, "product_id":2,"relation":"P", "user_name":"jane","product_name":"credit"}
                userproduct_3 = {"user_id":2, "product_id":3,"relation":"P", "user_name":"john","product_name":"indices"}
                userproduct_4 = {"user_id":1, "product_id":3,"relation":"S", "user_name":"jane","product_name":"indices"}
                userproduct_5 = {"user_id":2, "product_id":1,"relation":"S", "user_name":"john","product_name":"equities"}
                userproducts_to_add = [userproduct_1, userproduct_2, userproduct_3,userproduct_4, userproduct_5]

                with userproducts_table.batch_writer() as batch:
                    for item in userproducts_to_add:
                        response = batch.put_item(Item=item)   

            def handler(event, context):
                logger.debug(event)
                print(event)
                try:
                    _process(event, context, event['RequestType'])
                except Exception as e:
                    msg = "Failed to process custom resource event: {}".format(str(e))
                    logger.error(msg, exc_info=True)
                    response_body = {'Status': STATUS_FAILED, 'Data': {}, 'Reason': msg}
                    send(event, context, response_body["Status"],response_body["Data"])

      Description: A function to create sample data in the associated S3 buckets
      Handler: index.handler
      Role: !GetAtt SampleDataLambdaRole.Arn
      Runtime: python3.11
      Timeout: 300
      Environment:
        Variables:
          DataBucket: !Ref DataBucket
          UsersTableName: !GetAtt Users.Arn
          ProductsTableName: !GetAtt Products.Arn
          UserProductsTableName: !GetAtt UserProducts.Arn                  
# DynamoDB Tables
  Users:
    Type: AWS::DynamoDB::Table
    Properties:
      AttributeDefinitions:
        -
          AttributeName: "user_name"
          AttributeType: "S"         
      KeySchema:
        -
          AttributeName: "user_name"
          KeyType: "HASH"
      ProvisionedThroughput:
        ReadCapacityUnits: 5
        WriteCapacityUnits: 5
  Products:
    Type: AWS::DynamoDB::Table
    Properties:   
        AttributeDefinitions:
          -
            AttributeName: "product_id"
            AttributeType: "N" 
        KeySchema:
          -
            AttributeName: "product_id"
            KeyType: "HASH"
        ProvisionedThroughput:
          ReadCapacityUnits: 5
          WriteCapacityUnits: 5
  UserProducts:
    Type: AWS::DynamoDB::Table
    Properties:      
        AttributeDefinitions:
          -
            AttributeName: "user_name"
            AttributeType: "S"
          -
            AttributeName: "product_name"
            AttributeType: "S"

        KeySchema:
          -
            AttributeName: "user_name"
            KeyType: "HASH"
          -
            AttributeName: "product_name"
            KeyType: "RANGE"
        ProvisionedThroughput:
          ReadCapacityUnits: 5
          WriteCapacityUnits: 5
# Auth Lambda
  AuthLambda:
    Type: 'AWS::Lambda::Function'
    Properties:
      Code: 
          ZipFile: |
            import logging
            import os
            import boto3
            import json
            from boto3.dynamodb.conditions import Key

            data_bucket = os.environ['DataBucket']
            userRoleArn = os.environ['USER_ROLE']
            cognito_clientid = os.environ['CognitoClientId']

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

            client = boto3.client('cognito-idp')

            def handler(event, context):
                response = {}
                if (authenticated(event['username'], event['password'])):
                    username = event['username']
                    logger.debug("Auth success")
                    user_record = get_user_products(username)
                    scope_downpolicy = json.dumps(get_scope_downpolicy(username, user_record))
                    directory_mapping = json.dumps(get_directory_mapping(username, user_record))

                    response = {
                        'Role': userRoleArn,
                        'Policy': scope_downpolicy,
                        'HomeDirectoryType': "LOGICAL",
                        'HomeDirectoryDetails': directory_mapping
                    }

                logger.debug(response)
                return response


            def authenticated(username, password):
                authParameters = {
                    'USERNAME': username,
                    'PASSWORD': password
                }
                logger.debug('calling cognito')
                response = {}
                response = client.initiate_auth(AuthFlow='USER_PASSWORD_AUTH', AuthParameters=authParameters,
                                                ClientId=cognito_clientid)
                return response;


            def get_directory_mapping(username, user_record):
                dir_mapping = []
                user_root = {
                    "Entry": "/" + username,
                    "Target": "/" + data_bucket + "/" + username
                }
                dir_mapping.append(user_root)
                for product in user_record:
                    if product["relation"] == "P":
                        dir_mapping.append({
                            "Entry": "/published/" + product["product_name"],
                            "Target": "/" + data_bucket + "/" + product["product_name"]
                        })
                    else:
                        dir_mapping.append({
                            "Entry": "/subscribed/" + product["product_name"],
                            "Target": "/" + data_bucket + "/" + product["product_name"]
                        })
                return dir_mapping


            def get_scope_downpolicy(username, user_record):
                allow_listing = {
                    "Sid": "AllowListingOfFolder",
                    "Action": [
                        "s3:ListBucket"
                    ],
                    "Effect": "Allow",
                    "Resource": ["arn:aws:s3:::" + data_bucket],
                    "Condition": {
                        "StringLike": {
                            "s3:prefix": [
                            ]
                        }
                    }
                }
                allow_read = {
                    "Sid": "AllowObjectAccess",
                    "Effect": "Allow",
                    "Action": [
                        "s3:GetObject",
                        "s3:GetObjectVersion"
                    ],
                    "Resource": [

                    ]

                }
                allow_write = {
                    "Sid": "AllowObjectWrite",
                    "Effect": "Allow",
                    "Action": [
                        "s3:PutObject",
                        "s3:DeleteObject"
                    ],
                    "Resource": [

                    ]
                }
                policy = {
                    "Version": "2012-10-17",
                    "Statement": [
                    ]
                }

                file = "arn:aws:s3:::" + data_bucket + '/{product}/*'
                allow_listing['Condition']['StringLike']['s3:prefix'].append(username)
                allow_listing['Condition']['StringLike']['s3:prefix'].append(username + '/*')
                allow_read['Resource'].append(file.replace('{product}', username))
                allow_write['Resource'].append(file.replace('{product}', username))

                for product in user_record:
                    allow_listing['Condition']['StringLike']['s3:prefix'].append(product["product_name"])
                    allow_listing['Condition']['StringLike']['s3:prefix'].append(product["product_name"] + '/*')
                    allow_read['Resource'].append(file.replace('{product}', product["product_name"]))
                    if product["relation"] == "P":
                        allow_write['Resource'].append(file.replace('{product}', product["product_name"]))

                policy['Statement'].append(allow_listing)
                policy['Statement'].append(allow_read)
                policy['Statement'].append(allow_write)
                return policy;


            def get_user_products(user_name):
                userProductsTableName = os.environ['UserProductsTableName'].split('/')[1]
                dynamodb = boto3.resource('dynamodb') 
                userproducts_table = dynamodb.Table(userProductsTableName)
                query_results = userproducts_table.query(KeyConditionExpression=Key('user_name').eq(user_name))
                print(query_results)
                return query_results["Items"]

      Description: >-
        A function to provide IAM roles and policies for given user and
        serverId.
      Handler: index.handler
      Role: !GetAtt AuthLambdaRole.Arn
      Runtime: python3.11
      Timeout: 60
      Environment:
        Variables:
          DataBucket: !Ref DataBucket
          USER_ROLE: !GetAtt TransferUserRole.Arn 
          CognitoClientId: !Ref TransferUserPoolClient
          UserProductsTableName: !GetAtt UserProducts.Arn   
  AuthLambdaRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Statement:
          - Action:
              - 'sts:AssumeRole'
            Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
        Version: 2012-10-17
      ManagedPolicyArns:
        - 'arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
      Policies:
        - PolicyName: DynamoDBReadAccess
          PolicyDocument:
            Statement:
              - Action:
                  - "dynamodb:BatchGetItem"
                  - "dynamodb:ConditionCheckItem"
                  - "dynamodb:DescribeTable"
                  - "dynamodb:GetItem"
                  - "dynamodb:Scan"
                  - "dynamodb:Query"
                Effect: Allow
                Resource:
                  - !GetAtt Users.Arn
                  - !GetAtt Products.Arn
                  - !GetAtt UserProducts.Arn
            Version: 2012-10-17
  TransferUserPool:
      Type: AWS::Cognito::UserPool
      Properties:
        UsernameConfiguration: 
          CaseSensitive: false
        AutoVerifiedAttributes:
          - email
        Schema:
          - Name: email
            AttributeDataType: String
            Mutable: false
            Required: true
          - Name: name
            AttributeDataType: String
            Mutable: true
            Required: true
  TransferUserPoolClient:
    Type: AWS::Cognito::UserPoolClient
    Properties:
      UserPoolId: !Ref TransferUserPool
      ExplicitAuthFlows:
        - ALLOW_USER_PASSWORD_AUTH
        - ALLOW_REFRESH_TOKEN_AUTH
        - ALLOW_USER_SRP_AUTH
        - ALLOW_CUSTOM_AUTH
      AllowedOAuthFlowsUserPoolClient: true
      CallbackURLs:
        - http://localhost:3000
      AllowedOAuthFlows:
        - code
        - implicit
      AllowedOAuthScopes:
        - phone
        - email
        - openid
        - profile
      SupportedIdentityProviders:
        - COGNITO
Outputs:
  SftpIdpLambda:
    Description: Auth Lambda Function
    Value: !Ref AuthLambda
  CognitoUserPoolID:
    Value: !Ref TransferUserPool
    Description: Cognito UserPool ID
  CognitoAppClientID:
    Value: !Ref TransferUserPoolClient
    Description: Cognito App Client ID
  DataBucket:
    Value: !Ref DataBucket
    Description: S3 Bucket to store data for SFTP Server
