Exploring AWS Lambda code

Overview

There are two AWS Lambda functions that you deployed in the previous step. Both of them utilize the AWS SDK for Python (Boto3) library along with the Lambda Powertools Python via a Lambda layer to perform the Well-Architected Tool API access.

Deployed AWS Lambda functions

Click on each link to understand how each Lambda function works.

CreateWAWorkloadLambda.py

Overview

This Lambda function will create or update the workload and is idempotent based on the WorkloadName. The parameters for the workload are:

  • WorkloadName - The name of the workload. The name must be unique within an account within a Region. Spaces and capitalization are ignored when checking for uniqueness.
  • WorkloadDesc - The description for the workload.
  • WorkloadOwner - The review owner of the workload. The name, email address, or identifier for the primary group or individual that owns the workload review process.
  • WorkloadEnv - The environment for the workload. Valid Values: PRODUCTION | PREPRODUCTION
  • WorkloadRegion - The list of AWS Regions associated with the workload, for example, us-east-2, or ca-central-1. Maximum number of 50 items.
  • WorkloadLenses - The list of lenses associated with the workload. Each lens is identified by its LensAlias.
  • Tags - The tags to be associated with the workload. Maximum number of 50 items.

Python Code

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169

import botocore
import boto3
import json
import datetime
from aws_lambda_powertools import Logger
import jmespath
import cfnresponse
from pkg_resources import packaging

__author__    = "Eric Pullen"
__email__     = "eppullen@amazon.com"
__copyright__ = "Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved."
__credits__   = ["Eric Pullen"]

# Default region listed here
REGION_NAME = "us-east-1"
blankjson = {}
response = ""
logger = Logger()

# Helper class to convert a datetime item to JSON.
class DateTimeEncoder(json.JSONEncoder):
    def default(self, z):
        if isinstance(z, datetime.datetime):
            return (str(z))
        else:
            return super().default(z)

def CreateNewWorkload(
    waclient,
    workloadName,
    description,
    reviewOwner,
    environment,
    awsRegions,
    lenses,
    tags
    ):

    # Create your workload
    try:
        waclient.create_workload(
        WorkloadName=workloadName,
        Description=description,
        ReviewOwner=reviewOwner,
        Environment=environment,
        AwsRegions=awsRegions,
        Lenses=lenses,
        Tags=tags
        )
    except waclient.exceptions.ConflictException as e:
        workloadId,workloadARN = FindWorkload(waclient,workloadName)
        logger.warning("WARNING - The workload name %s already exists as workloadId %s" % (workloadName, workloadId))
        UpdateWorkload(waclient,workloadId,workloadARN, workloadName,description,reviewOwner,environment,awsRegions,lenses,tags)

        # Maybe we should "update" the above variables?

    except botocore.exceptions.ParamValidationError as e:
        logger.error("ERROR - Parameter validation error: %s" % e)
    except botocore.exceptions.ClientError as e:
        logger.error("ERROR - Unexpected error: %s" % e)


def FindWorkload(
    waclient,
    workloadName
    ):

    # Finding your WorkloadId
    try:
        response=waclient.list_workloads(
        WorkloadNamePrefix=workloadName
        )
    except botocore.exceptions.ParamValidationError as e:
        logger.error("ERROR - Parameter validation error: %s" % e)
    except botocore.exceptions.ClientError as e:
        logger.error("ERROR - Unexpected error: %s" % e)

    # print("Full JSON:",json.dumps(response['WorkloadSummaries'], cls=DateTimeEncoder))
    workloadId = response['WorkloadSummaries'][0]['WorkloadId']
    workloadARN = response['WorkloadSummaries'][0]['WorkloadArn']
    # print("WorkloadId",workloadId)
    return workloadId, workloadARN

def UpdateWorkload(
    waclient,
    workloadId,
    workloadARN,
    workloadName,
    description,
    reviewOwner,
    environment,
    awsRegions,
    lenses,
    tags
    ):

    logger.info("Updating workload properties")
    # Create your workload
    try:
        waclient.update_workload(
        WorkloadId=workloadId,
        WorkloadName=workloadName,
        Description=description,
        ReviewOwner=reviewOwner,
        Environment=environment,
        AwsRegions=awsRegions,
        )
    except botocore.exceptions.ParamValidationError as e:
        logger.error("ERROR - Parameter validation error: %s" % e)
    except botocore.exceptions.ClientError as e:
        logger.error("ERROR - Unexpected error: %s" % e)
    # Should add updates for the lenses?
    # Should add the tags as well
    logger.info("Updating workload tags")
    try:
        waclient.tag_resource(WorkloadArn=workloadARN,Tags=tags)
    except botocore.exceptions.ParamValidationError as e:
        logger.error("ERROR - Parameter validation error: %s" % e)
    except botocore.exceptions.ClientError as e:
        logger.error("ERROR - Unexpected error: %s" % e)



def lambda_handler(event, context):
    boto3_min_version = "1.16.38"
    # Verify if the version of Boto3 we are running has the wellarchitected APIs included
    if (packaging.version.parse(boto3.__version__) < packaging.version.parse(boto3_min_version)):
        logger.error("Your Boto3 version (%s) is less than %s. You must ugprade to run this script (pip3 upgrade boto3)" % (boto3.__version__, boto3_min_version))
        exit()
    responseData = {}
    # print(json.dumps(event))
    try:
        WORKLOADNAME = event['ResourceProperties']['WorkloadName']
        DESCRIPTION = event['ResourceProperties']['WorkloadDesc']
        REVIEWOWNER = event['ResourceProperties']['WorkloadOwner']
        ENVIRONMENT= event['ResourceProperties']['WorkloadEnv']
        AWSREGIONS = [event['ResourceProperties']['WorkloadRegion']]
        LENSES = event['ResourceProperties']['WorkloadLenses']
        TAGS = event['ResourceProperties']['Tags']
        SERVICETOKEN = event['ResourceProperties']['ServiceToken']
    except:
        responseData['Error'] = "ERROR LOADING RESOURCE PROPERTIES"
        cfnresponse.send(event, context, cfnresponse.FAILED, responseData, 'createWAWorkloadHelperFunction')
        exit()

    IncomingARN = SERVICETOKEN.split(":")
    REGION_NAME = IncomingARN[3]


    logger.info("Starting Boto %s Session in %s" % (boto3.__version__, REGION_NAME))
    # Create a new boto3 session
    SESSION = boto3.session.Session()
    # Initiate the well-architected session using the region defined above
    WACLIENT = SESSION.client(
        service_name='wellarchitected',
        region_name=REGION_NAME,
    )

    logger.info("Creating a new workload")
    CreateNewWorkload(WACLIENT,WORKLOADNAME,DESCRIPTION,REVIEWOWNER,ENVIRONMENT,AWSREGIONS,LENSES,TAGS)
    logger.info("Finding your WorkloadId")
    workloadId,workloadARN = FindWorkload(WACLIENT,WORKLOADNAME)
    logger.info("New workload created with id %s" % workloadId)
    responseData['WorkloadId'] = workloadId
    responseData['WorkloadARN'] = workloadARN
    logger.info("Response will be %s" % responseData)

    cfnresponse.send(event, context, cfnresponse.SUCCESS, responseData)

Example incoming Lambda event

This is an Lambda test event you can use to see if the Lambda function works as expected:

{
  "RequestType": "Create",
  "ResponseURL": "http://pre-signed-S3-url-for-response",
  "StackId": "arn:aws:cloudformation:us-east-1:123456789012:stack/MyStack/guid",
  "RequestId": "unique id for this create request",
  "ResourceType": "Custom::CreateNewWAFRFunction",
  "LogicalResourceId": "CreateNewWAFRFunction",
  "ResourceProperties": {
    "ServiceToken": "arn:aws:lambda:us-east-1:123456789012:function:CreateNewWAFRFunction",
    "WorkloadName": "Lambda WA Workload Test",
    "WorkloadOwner": "Lambda Script",
    "WorkloadDesc": "Test Lambda WA Workload",
    "WorkloadRegion": "us-east-1",
    "WorkloadLenses": [
      "wellarchitected",
      "serverless"
    ],
    "WorkloadEnv": "PRODUCTION",
    "Tags": {
      "TestTag3": "TestTagValue4",
      "TestTag": "TestTagValue"
    }
  }
}

IAM Privileges

Using the concept of least privileges for each AWS Lambda function, we create an IAM role for this function that only allows certain access to the Well-Architected Tool (CreateWorkload and UpdateWorkload specifically) as well as the ability to create log file entries (lines 29-34).

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

CreateWAlambdaIAMRole:
  Type: AWS::IAM::Role
  Properties:
    AssumeRolePolicyDocument:
      Version: 2012-10-17
      Statement:
        - Action:
            - sts:AssumeRole
          Effect: Allow
          Principal:
            Service:
              - lambda.amazonaws.com
    Policies:
      - PolicyDocument:
          Version: 2012-10-17
          Statement:
            - Action:
                - logs:CreateLogGroup
                - logs:CreateLogStream
                - logs:PutLogEvents
              Effect: Allow
              Resource:
                - !Sub arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/${CreateWALambdaFunctionName}:*
        PolicyName: lambda
      - PolicyDocument:
          Version: 2012-10-17
          Statement:
            - Action:
                - wellarchitected:CreateWorkload
                - wellarchitected:GetWorkload
                - wellarchitected:List*
                - wellarchitected:TagResource
                - wellarchitected:UntagResource
                - wellarchitected:UpdateWorkload
              Effect: Allow
              Resource: '*'
        PolicyName: watool

UpdateWAQuestionLambda.py

Overview

This Lambda function will update the answer to a specific question in a workload review. The parameters for the workload are:

  • WorkloadId - The ID assigned to the workload. This ID is unique within an AWS Region.
  • Lens - The alias of the lens, for example, wellarchitected or serverless.
  • Pillar - The ID used to identify a pillar, for example, security.
  • QuestionAnswers - An array of pillar Questions and associated best practices you wish to mark as selected.

Python Code

import botocore
import boto3
import json
import datetime
from aws_lambda_powertools import Logger
import jmespath
import cfnresponse
from pkg_resources import packaging

__author__    = "Eric Pullen"
__email__     = "eppullen@amazon.com"
__copyright__ = "Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved."
__credits__   = ["Eric Pullen"]

# Default region listed here
REGION_NAME = "us-east-1"
blankjson = {}
response = ""
logger = Logger()

# Helper class to convert a datetime item to JSON.
class DateTimeEncoder(json.JSONEncoder):
    def default(self, z):
        if isinstance(z, datetime.datetime):
            return (str(z))
        else:
            return super().default(z)

def findQuestionId(
    waclient,
    workloadId,
    lensAlias,
    pillarId,
    questionTitle
    ):

    # Find a questionID using the questionTitle
    try:
        response=waclient.list_answers(
        WorkloadId=workloadId,
        LensAlias=lensAlias,
        PillarId=pillarId
        )
    except botocore.exceptions.ParamValidationError as e:
        logger.error("ERROR - Parameter validation error: %s" % e)
    except botocore.exceptions.ClientError as e:
        logger.error("ERROR - Unexpected error: %s" % e)

    answers = response['AnswerSummaries']
    while "NextToken" in response:
        response = waclient.list_answers(WorkloadId=workloadId,LensAlias=lensAlias,PillarId=pillarId,NextToken=response["NextToken"])
        answers.extend(response["AnswerSummaries"])

    jmesquery = "[?starts_with(QuestionTitle, `"+questionTitle+"`) == `true`].QuestionId"
    questionId = jmespath.search(jmesquery, answers)

    return questionId[0]

def findChoiceId(
    waclient,
    workloadId,
    lensAlias,
    questionId,
    choiceTitle,
    ):

    # Find a choiceId using the choiceTitle
    try:
        response=waclient.get_answer(
        WorkloadId=workloadId,
        LensAlias=lensAlias,
        QuestionId=questionId
        )
    except botocore.exceptions.ParamValidationError as e:
        logger.error("ERROR - Parameter validation error: %s" % e)
    except botocore.exceptions.ClientError as e:
        logger.error("ERROR - Unexpected error: %s" % e)

    jmesquery = "Answer.Choices[?starts_with(Title, `"+choiceTitle+"`) == `true`].ChoiceId"
    choiceId = jmespath.search(jmesquery, response)

    return choiceId[0]

def updateAnswersForQuestion(
    waclient,
    workloadId,
    lensAlias,
    questionId,
    selectedChoices,
    notes
    ):

    # Update a answer to a question
    try:
        response=waclient.update_answer(
        WorkloadId=workloadId,
        LensAlias=lensAlias,
        QuestionId=questionId,
        SelectedChoices=selectedChoices,
        Notes=notes
        )
    except botocore.exceptions.ParamValidationError as e:
        logger.error("ERROR - Parameter validation error: %s" % e)
    except botocore.exceptions.ClientError as e:
        logger.error("ERROR - Unexpected error: %s" % e)

    # print(json.dumps(response))
    jmesquery = "Answer.SelectedChoices"
    answers = jmespath.search(jmesquery, response)
    return answers

def lambda_handler(event, context):
    boto3_min_version = "1.16.38"
    # Verify if the version of Boto3 we are running has the wellarchitected APIs included
    if (packaging.version.parse(boto3.__version__) < packaging.version.parse(boto3_min_version)):
        logger.error("Your Boto3 version (%s) is less than %s. You must ugprade to run this script (pip3 upgrade boto3)" % (boto3.__version__, boto3_min_version))
        exit()
    responseData = {}
    print(json.dumps(event))
    try:
        WORKLOADID = event['ResourceProperties']['WorkloadId']
        PILLAR = event['ResourceProperties']['Pillar']
        LENS = event['ResourceProperties']['Lens']
        QUESTIONANSWERS = event['ResourceProperties']['QuestionAnswers']
        SERVICETOKEN = event['ResourceProperties']['ServiceToken']
    except:
        responseData['Error'] = "ERROR LOADING RESOURCE PROPERTIES"
        cfnresponse.send(event, context, cfnresponse.FAILED, responseData, 'createWAWorkloadHelperFunction')
        exit()
    IncomingARN = SERVICETOKEN.split(":")
    REGION_NAME = IncomingARN[3]


    logger.info("Starting Boto %s Session in %s" % (boto3.__version__, REGION_NAME))
    # Create a new boto3 session
    SESSION = boto3.session.Session()
    # Initiate the well-architected session using the region defined above
    WACLIENT = SESSION.client(
        service_name='wellarchitected',
        region_name=REGION_NAME,
    )

    for qaList in QUESTIONANSWERS:
        for question, answerList in qaList.items():
            print(question, answerList)
            # First we must find the questionID
            questionId = findQuestionId(WACLIENT,WORKLOADID,LENS,PILLAR,question)
            logger.info("Found QuestionID of '%s' for the question text of '%s'" % (questionId, question))
            choiceSet = []
            # Now we build the choice selection based on the answers provided
            for answers in answerList:
                choiceSet.append(findChoiceId(WACLIENT,WORKLOADID,LENS,questionId,answers))
            logger.info("All choices we will select for questionId of %s is %s" % (questionId, choiceSet))
            # Update the answer for the question
            updateAnswersForQuestion(WACLIENT,WORKLOADID,LENS,questionId,choiceSet,'Added by Python')
    # exit()
    cfnresponse.send(event, context, cfnresponse.SUCCESS, responseData, 'createWAWorkloadHelperFunction')

Example incoming Lambda event

This is an Lambda test event you can use to see if the Lambda function works as expected:

{
  "RequestType": "Create",
  "ServiceToken": "arn:aws:lambda:us-east-1:123456789012:function:UpdateWAQFunction",
  "ResponseURL": "http://pre-signed-S3-url-for-response",
  "StackId": "arn:aws:cloudformation:us-east-1:123456789012:stack/MyStack/guid",
  "RequestId": "unique id for this create request",
  "LogicalResourceId": "UpdateWAQFunction",
  "ResourceType": "Custom::UpdateWAQFunction",
  "ResourceProperties": {
    "ServiceToken": "arn:aws:lambda:us-east-1:123456789012:function:UpdateWAQFunction",
    "QuestionAnswers": [
      {
        "How do you determine what your priorities are": [
          "Evaluate governance requirements",
          "Evaluate compliance requirements"
        ]
      },
      {
        "How do you reduce defects, ease remediation, and improve flow into production": [
          "Use version control",
          "Perform patch management",
          "Use multiple environments"
        ]
      }
    ],
    "Pillar": "operationalExcellence",
    "Lens": "wellarchitected",
    "WorkloadId": "d1a1d1a1d1a1d1a1d1a1d1a1d1a1d1a1"
  }
}

IAM Privileges

Using the concept of least privileges for each AWS Lambda function, we create an IAM role for this function that only allows certain access to the Well-Architected Tool (UpdateAnswer specifically) as well as the ability to create log file entries (lines 29-34).

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

UpdateWAQlambdaIAMRole:
  Type: AWS::IAM::Role
  Properties:
    AssumeRolePolicyDocument:
      Version: 2012-10-17
      Statement:
        - Action:
            - sts:AssumeRole
          Effect: Allow
          Principal:
            Service:
              - lambda.amazonaws.com
    Policies:
      - PolicyDocument:
          Version: 2012-10-17
          Statement:
            - Action:
                - logs:CreateLogGroup
                - logs:CreateLogStream
                - logs:PutLogEvents
              Effect: Allow
              Resource:
                - !Sub arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/${UpdateWAQLambdaFunctionName}:*
        PolicyName: lambda
      - PolicyDocument:
          Version: 2012-10-17
          Statement:
            - Action:
                - wellarchitected:GetAnswer
                - wellarchitected:GetWorkload
                - wellarchitected:List*
                - wellarchitected:TagResource
                - wellarchitected:UntagResource
                - wellarchitected:UpdateAnswer
              Effect: Allow
              Resource: '*'
        PolicyName: watool