# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

import boto3
import os
import json
from botocore.exceptions import ClientError
from botocore.config import Config

# general env variables
bucket_name = os.environ['sourceBucketName']
allow_origins = os.environ['allowOrigins']
region = os.environ['region']
s3_region_endpoint = f'https://s3.{region}.amazonaws.com'
s3_department_tag_key = os.environ['s3DepartmentTagKey']

# Cognito env variables
id_pool_id = os.environ['idPoolId']
user_pool_id = os.environ['userPoolId']
id_login_provider = f'cognito-idp.{region}.amazonaws.com/{user_pool_id}'

# create Cognito client
id_client = boto3.client('cognito-identity')

def lambda_handler(event, context):

    body = json.loads(event['body'])
    fileName = body['fileName']
    fileType = body['fileType']

    # get the preferred cognito group name and arn
    id_token_claims = event['requestContext']['authorizer']['claims']
    preferred_group_name = id_token_claims['department']
    preferred_role_arn = id_token_claims['cognito:preferred_role']

    # get the id token from the request header
    id_token = event['headers']['Authorization']

    # get Cognito id from the Identity Pool
    identity_response = id_client.get_id(
        IdentityPoolId = id_pool_id,
        Logins = {id_login_provider: id_token})

    # get temporary AWS credentials from the Identity Pool based on Cognito id
    identity_cred=id_client.get_credentials_for_identity(
        CustomRoleArn = preferred_role_arn,
        IdentityId = identity_response['IdentityId'],
        Logins = {id_login_provider: id_token})

    temp_aws_credentials=identity_cred["Credentials"]

    # create the s3 client
    s3_client=boto3.client(
        's3',
        endpoint_url = s3_region_endpoint,
        aws_access_key_id = temp_aws_credentials['AccessKeyId'],
        aws_secret_access_key = temp_aws_credentials['SecretKey'],
        aws_session_token = temp_aws_credentials['SessionToken'],
        region_name = region,
        config = Config(signature_version='s3v4'))

    # create presigned s3 url to upload the object and tag it for downstream access control
    try:
        params={'Bucket': bucket_name, 'Key': preferred_group_name+'/'+fileName,
                  'ContentType': fileType, 'Tagging': '{0}={1}'.format(s3_department_tag_key,preferred_group_name)}
        presignedurl = s3_client.generate_presigned_url(
            'put_object',
            params
        )
    except ClientError as error:
        print(error)
        return {
            'statusCode': 500,
            'headers': {
                'Access-Control-Allow-Origin': allow_origins,
                'Access-Control-Allow-Credentials': True
            }
        }

    # return the presigned url and the preferred group name for tagging the object
    return {
        'statusCode': 200,
        'headers': {
            'Access-Control-Allow-Origin': allow_origins,
            'Access-Control-Allow-Credentials': True
        },
        'body': json.dumps({'preSignedUrl': presignedurl, 'tagKey': s3_department_tag_key, 'tagValue': preferred_group_name})
    }