######################## Apache License ########################################

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.

# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
# the Software, and to permit persons to whom the Software is furnished to do so.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.


#################################################################################


#cloudformation template which creates S3 bucket and bedrock knowledgebase
AWSTemplateFormatVersion: "2010-09-09"
Description: Cloud Formation template to deploy resources for supply chain blog

################# parameter groups ###############################

Metadata:
   AWS::CloudFormation::Interface:
    ParameterGroups:
      - Label:
          default: "Email Configuration"
        Parameters:
          - Fromemailaddress
          - Toemailaddress
      - Label:
          default: "VPC configurations"
          Parameters:
          - VPCID
          - VPCSubnets
          - VPCSecurityGroupIds
          

################ parameters ####################################         

Parameters:


  # Email Configuration
  Toemailaddress:
    Type: String
    Default: "abc@gmail.com"
    AllowedPattern: '[^@]+@[^@]+\.[^@]+'
    Description: Email address to receive email
   
  Fromemailaddress:
    Type: String
    Default: "xyz@gmail.com"
    AllowedPattern: '[^@]+@[^@]+\.[^@]+'
    Description: Email address to send email

  VPCId: 
    Type: AWS::EC2::VPC::Id
    Description: "VPCId in which the lambda is to be deployed"


  VPCSubnets: 
      Description : "SubnetId in VPC"
      Type : List<AWS::EC2::Subnet::Id>
      ConstraintDescription : "Must be a valid subnet ID"


  # add security group id
  VPCSecurityGroupIds:
    Type: List<AWS::EC2::SecurityGroup::Id>
    Description: Security group(s) for access

Resources:

############### KMS Key for S3 bucket ###################
  MyKMSKey:
    Type: AWS::KMS::Key
    Properties:
      Description: "Encryption key for S3 bucket"
      EnableKeyRotation: true
      KeyPolicy:
        Version: "2012-10-17"
        Id: key-default-1
        Statement:
          - Sid: Enable IAM User Permissions
            Effect: Allow
            Principal:
              AWS: !Join 
                - ''
                - - 'arn:aws:iam::'
                  - !Ref 'AWS::AccountId'
                  - ':root'
            Action: 'kms:*'
            Resource: '*'
          - Sid: VPC Flow Logs Access
            Effect: Allow
            Principal:
              Service: logs.amazonaws.com
            Action:
              - kms:Encrypt
              - kms:Decrypt
              - kms:ReEncrypt*
              - kms:GenerateDataKey*
              - kms:DescribeKey
            Resource: '*'
            Condition:
              StringEquals:
                "aws:SourceAccount": !Ref 'AWS::AccountId'
              ArnLike:
                "aws:SourceArn": !Sub "arn:aws:logs:${AWS::Region}:${AWS::AccountId}:*" 


#############################  S3 buckets ################################


  S3AccessLogsBucket:
    Type: AWS::S3::Bucket
    Properties:
      PublicAccessBlockConfiguration:
        BlockPublicPolicy: true
        RestrictPublicBuckets: true
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: 'aws:kms'
              KMSMasterKeyID: !GetAtt MyKMSKey.Arn
      VersioningConfiguration:
        Status: Enabled
      Tags:
        - Key: Purpose
          Value: S3AccessLogs
        - Key: ComplianceException
          Value: LoggingBucket


  S3AccessLogsBucketPolicy:
      Type: AWS::S3::BucketPolicy
      Properties:
        Bucket: !Ref S3AccessLogsBucket
        PolicyDocument:
          Version: '2012-10-17'
          Statement:
            - Sid: S3ServerAccessLogsPolicy
              Effect: Allow
              Principal:
                Service: logging.s3.amazonaws.com
              Action: s3:PutObject
              Resource: !Sub '${S3AccessLogsBucket.Arn}/*'
              Condition:
                ArnLike:
                  aws:SourceArn: !Sub 'arn:aws:s3:::${BedrockKBQDataSourceBucket}'
                StringEquals:
                  aws:SourceAccount: !Ref 'AWS::AccountId'

  BedrockKBQDataSourceBucket:
      Type: AWS::S3::Bucket
      Properties:
        PublicAccessBlockConfiguration:
          BlockPublicPolicy: true
          RestrictPublicBuckets: true
        ObjectLockEnabled: true
        BucketEncryption:
          ServerSideEncryptionConfiguration:
            - ServerSideEncryptionByDefault:
                SSEAlgorithm: 'aws:kms'
                KMSMasterKeyID: !GetAtt MyKMSKey.Arn
        LoggingConfiguration:
          DestinationBucketName: !Ref S3AccessLogsBucket
          LogFilePrefix: 'data-source-bucket-logs/'
        VersioningConfiguration:
          Status: Enabled

  BedrockKBDSBucketPolicy:
    Type: AWS::S3::BucketPolicy
    Properties:
      Bucket: !Ref BedrockKBQDataSourceBucket
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Sid: EnforceTLSRequestsOnly
            Effect: Deny
            Principal: '*'
            Action: 's3:*'
            Resource:
              - !Sub '${BedrockKBQDataSourceBucket.Arn}/*'
              - !GetAtt BedrockKBQDataSourceBucket.Arn
            Condition:
              Bool:
                'aws:SecureTransport': false
          - Sid: EnforceEncryptionInTransit
            Effect: Deny
            Principal: '*'
            Action: 's3:*'
            Resource:
              - !Sub '${BedrockKBQDataSourceBucket.Arn}/*'
              - !GetAtt BedrockKBQDataSourceBucket.Arn
            Condition:
              NumericLessThan:
                's3:TlsVersion': 1.2
          - Sid: DenyIncorrectEncryptionHeader
            Effect: Deny
            Principal: '*'
            Action: 's3:PutObject'
            Resource: !Sub '${BedrockKBQDataSourceBucket.Arn}/*'
            Condition:
              StringNotEquals:
                's3:x-amz-server-side-encryption-aws-kms-key-id': !GetAtt MyKMSKey.Arn

  
 # S3 bucket for storing Device 2 mitigations
  Device2MitigationsBucket:
      Type: AWS::S3::Bucket
      Properties:
        PublicAccessBlockConfiguration:
          BlockPublicPolicy: true
          RestrictPublicBuckets: true
        ObjectLockEnabled: true
        BucketEncryption:
          ServerSideEncryptionConfiguration:
            - ServerSideEncryptionByDefault:
                SSEAlgorithm: 'aws:kms'
        LoggingConfiguration:
          DestinationBucketName: !Ref S3AccessLogsBucket
          LogFilePrefix: 'device2-mitigations-bucket-logs/'
        VersioningConfiguration:
          Status: Enabled

  Device2MitigationsBucketPolicy:
    Type: AWS::S3::BucketPolicy
    Properties:
      Bucket: !Ref Device2MitigationsBucket
      PolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Sid: EnforceTLSRequestsOnly
            Effect: Deny
            Principal: '*'
            Action: 's3:*'
            Resource:
              - !Sub '${Device2MitigationsBucket.Arn}/*'
              - !GetAtt Device2MitigationsBucket.Arn
            Condition:
              Bool:
                'aws:SecureTransport': false
          - Sid: EnforceEncryptionInTransit
            Effect: Deny
            Principal: '*'
            Action: 's3:*'
            Resource:
              - !Sub '${Device2MitigationsBucket.Arn}/*'
              - !GetAtt Device2MitigationsBucket.Arn
            Condition:
              NumericLessThan:
                's3:TlsVersion': 1.2
          - Sid: DenyIncorrectEncryptionHeader
            Effect: Deny
            Principal: '*'
            Action: 's3:PutObject'
            Resource: !Sub '${Device2MitigationsBucket.Arn}/*'
            Condition:
              StringNotEquals:
                's3:x-amz-server-side-encryption-aws-kms-key-id': !GetAtt MyKMSKey.Arn

  

########################## Amazon SES resources ####################################


  ToEmail:
    Type: AWS::SES::EmailIdentity
    Properties:
      EmailIdentity: !Ref Toemailaddress

  FromEmail:
    Type: AWS::SES::EmailIdentity
    Properties:
      EmailIdentity: !Ref Fromemailaddress


############################# DynamoDB resources #################################

  DeviceClassification:
    Type: AWS::DynamoDB::Table
    DeletionPolicy: Delete
    Properties:
      BillingMode: PROVISIONED
      AttributeDefinitions:
        -
          AttributeName: "device_name"
          AttributeType: "S"
        -
          AttributeName: "classification"
          AttributeType: "N"  
      KeySchema:
        - AttributeName: device_name
          KeyType: HASH
        - AttributeName: classification
          KeyType: RANGE
      ProvisionedThroughput:
        ReadCapacityUnits: 5
        WriteCapacityUnits: 5
      # Add SSE Configuration
      SSESpecification:
        SSEEnabled: true
        # Optional: Specify KMS key
        SSEType: KMS
        KMSMasterKeyId: alias/aws/dynamodb
      PointInTimeRecoverySpecification:
        PointInTimeRecoveryEnabled: true

  
# role for lambda function which populates dynamodb table
  LambdaRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Principal:
            Service:
            - lambda.amazonaws.com
          Action:
          - sts:AssumeRole
      Path: "/"
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
        - arn:aws:iam::aws:policy/service-role/AWSLambdaVPCAccessExecutionRole
      Policies:
        - PolicyName: dynamodbAccessPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
            - Effect: Allow
              Action:
                - dynamodb:PutItem
                - dynamodb:BatchWriteItem
                - dynamodb:GetItem
                - dynamodb:BatchGetItem
                - dynamodb:Scan
                - dynamodb:Query
                - dynamodb:UpdateItem
                - dynamodb:DescribeTable
              Resource: 
                - !GetAtt DeviceClassification.Arn
                - !Sub '${DeviceClassification.Arn}/index/*'
        - PolicyName: LogsAccessPolicyName
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
            - Effect: Allow
              Action:
                - logs:CreateLogGroup
                - logs:CreateLogStream
                - logs:PutLogEvents
              Resource: 
                - !Sub 'arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/${AWS::StackName}-*'


# Lambda function which populates the dynamodb table     
  InitializeTableFunction:
    Type: AWS::Lambda::Function
    Properties:
    # add lambda code to add values to dynamodb
      Handler: index.lambda_handler
      Code: 
        ZipFile: |
          import json
          import boto3
          import cfnresponse
          import os

          TABLE_NAME = os.environ.get('TABLE_NAME')
          REGION = os.environ.get('AWS_REGION')

          def lambda_handler(event, context):
              # Initialize responseData dictionary at the beginning
              responseData = {}
              
              try:
                  dynamodb = boto3.resource('dynamodb', region_name=REGION)
                  table = dynamodb.Table(TABLE_NAME)

                  # add multiple items to dynamodb routeTableName:
                 
                  with table.batch_writer() as batch:
                    batch.put_item(
                        Item={
                            'device_name': 'HIV diagnostic tests',
                            'classification': 3

                        }
                    )
                    batch.put_item(
                        Item={
                            'device_name': 'syringes',
                            'classification': 1
                        }
                    )   
                    batch.put_item(
                        Item={
                            'device_name': 'blood pressure cuff',
                            'classification': 2
                        }
                    )
                    batch.put_item(
                        Item={
                            'device_name': 'cochlear implants',
                            'classification': 3
                        }
                    )              
                  
                  response = table.put_item(
                      Item={
                          'device_name': 'oxygen mask',
                          'classification': 2
                      }
                  )
                  
                  # Add successful response data
                  responseData['Data'] = response
                  responseData['Message'] = 'Item added successfully'
                  
                  # Send success response to CloudFormation
                  cfnresponse.send(event, context, cfnresponse.SUCCESS, responseData)
                  
                  return {
                      'statusCode': 200,
                      'body': json.dumps(responseData)
                  }
                  
              except Exception as e:
                  print(f"Error: {str(e)}")
                  
                  # Add error information to responseData
                  responseData['Error'] = str(e)
                  responseData['Message'] = 'Error adding item to the database'
                  
                  # Send failure response to CloudFormation
                  cfnresponse.send(event, context, cfnresponse.FAILED, responseData)
                  
                  return {
                      'statusCode': 500,
                      'body': json.dumps(responseData)
                  }
              
      Role: !GetAtt LambdaRole.Arn
      Runtime: python3.12
      Environment:
        Variables:
          TABLE_NAME: !Ref DeviceClassification
      Timeout: 60
      ReservedConcurrentExecutions: 2

  InitializeDynamoDB:
      Type: Custom::InitializeTable
      Properties:
        ServiceToken: !GetAtt InitializeTableFunction.Arn
        TableName: !Ref DeviceClassification
        Version: "1.0"  # Add version to trigger updates when needed
  

  
 ############################ Lamba functions ##########################

  #supplychain lambda executionrole
  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
        - arn:aws:iam::aws:policy/AmazonDynamoDBReadOnlyAccess
        - arn:aws:iam::aws:policy/service-role/AWSLambdaVPCAccessExecutionRole

  #supplychain Lamba fucntion      
  SupplyChainMgmtFunction:
    Type: AWS::Lambda::Function
    Properties:
      Description: A Lambda function to finds the classification of a device by querying dynamodb
      Handler: lambda_function.lambda_handler
      Role: !GetAtt LambdaExecutionRole.Arn
      VpcConfig:
        SecurityGroupIds: !Ref VPCSecurityGroupIds
        SubnetIds: !Ref VPCSubnets
      Code:
        S3Bucket: aws-blogs-artifacts-public
        S3Key: ML-17903/Lambda/supplychainmgmt.zip
      Runtime: python3.12
      Timeout: 30
      MemorySize: 128   
      ReservedConcurrentExecutions: 100
      Environment:
        Variables:
          TABLE_NAME: !Ref DeviceClassification
          DYNAMODB_ENDPOINT: !Sub 'dynamodb.${AWS::Region}.amazonaws.com'
    

# add permissions for bedrock to invokelambda
  InvokePermissionDynamodbLambda:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !GetAtt SupplyChainMgmtFunction.Arn
      Action: lambda:InvokeFunction
      Principal: bedrock.amazonaws.com    



  # SES Lambda execution role
  SESLambdaExecutionRole:
      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
          - arn:aws:iam::aws:policy/service-role/AWSLambdaVPCAccessExecutionRole
        #policy to read from s3 bucket
        Policies:
          - PolicyName: SESAccess
            PolicyDocument:
              Version: '2012-10-17'
              Statement:
                - Effect: Allow
                  Action: 
                    - ses:SendEmail
                    - ses:SendRawEmail
                    - ses:SendTemplatedEmail
                  Resource: !Sub 'arn:aws:ses:${AWS::Region}:${AWS::AccountId}:*'       

  # SES Lambda execution role
  EmailReviewersFunction:
    Type: AWS::Lambda::Function
    Properties:
      Description: A Lambda function to check objects in S3
      Handler: lambda_function.lambda_handler
      Role: !GetAtt SESLambdaExecutionRole.Arn
      VpcConfig:
        SecurityGroupIds: !Ref VPCSecurityGroupIds
        SubnetIds: !Ref VPCSubnets
      Code:
        S3Bucket: aws-blogs-artifacts-public
        S3Key: ML-17903/Lambda/emailreviewers.zip
      Runtime: python3.12
      Timeout: 30
      MemorySize: 128 
      ReservedConcurrentExecutions: 100
      Environment:
        Variables:
          FROM_EMAIL: !Ref Fromemailaddress
          

  # add permissions for bedrock to invokelambda
  InvokePermissionSESLambda:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !GetAtt EmailReviewersFunction.Arn
      Action: lambda:InvokeFunction
      Principal: bedrock.amazonaws.com  


#######################  Outputs #######################

Outputs:
  Device2MitigationsBucket:
    Value: !Ref Device2MitigationsBucket
    Export:
      Name: Device2MitigationsBucket
  
  BedrockKBQDataSourceBucket:
    Value: !Ref BedrockKBQDataSourceBucket
    Export:
      Name: BedrockKBQDataSourceBucket
      # add export for dynamodb table    # add export for dynamodb table
  DeviceClassification:
    Value: !GetAtt DeviceClassification.Arn
    Export:
      Name: DeviceClassification
      # add export for lambda function
  SupplyChainMgmtFunction:
    Value: !GetAtt SupplyChainMgmtFunction.Arn
    Export:
      Name: DeviceClassificationFunction
      # add export for lambda function
  EmailReviewersFunction:
    Value: !GetAtt EmailReviewersFunction.Arn
    Export:
      Name: EmailReviewersFunction



          
