【Workaround】How to use the latest Boto3 with SSM Automation (or Lambda)

Created date: 2024/05/19, Update date: 2024/05/19

SSM Automation aws:executeScript may not be able to run the latest API.
For example, the Basseline API of Control Tower is now available on December 15, 2024, but at the time of writing this article, it is not available on aws:executeScript.
SSM Automation aws:executeScript action and Step Functions are the same, and the latest APIs may not be available.
In order to respond to this, we have examined the version of Boto3 in the SSM Automation aws:executeScript action and used the latest API.
Simply upgrading Boto3 with pip did not work, so I will describe the corresponding method.

Table of contents

What is Automation?
First Conclusion
1. In the case of SSM Automation
2. For Lambda
Confirmation/verification result (for SSM Automation)
SSM Automation Terraform code used for verification
Terraform code used to verify the version upgrade of Boto3 of Lambda

What is Automation?

Automation is one of the functions of AWS Systems Manager and can automate development operation tasks.
Various operations can be automated by using Rambook, which defines the tasks executed by automation.
There are many Rambooks owned by AWS, so you can execute the task to prepare by passing the necessary parameters for the user side.
- Reference: Systems Manager Automation runbook reference - AWS Systems Manager Automation runbook reference

You can also define your own run books, implement your own processing by combining actions, or call multiple Rambooks owned by AWS to form an operation task flow.
- Reference: Systems Manager Automation actions reference - AWS Systems Manager

First Conclusion

In the case of SSM Automation

As a prerequisite, you need to specify a zip file including the package file of the boto3 / botocore in the Attachment of SSM Automation.
You can use the latest Boto3 with the following two processing:

AWS_DATA_PATH Specify a directory ( /botocore/data placed under the current directory distribution) where the latest model file is stored in the environment variable.

The directory that runs SSM Automation is a directory with a random name, and you can use the latest model file by specifying /botocore/data placed under the current directory.


          import os

          os.environ["AWS_DATA_PATH"] = f"{os.getcwd()}/botocore/data"

Reload boto3, botocore.


          import os
          import sys
          import boto3
          import botocore

          sys.path.insert(0, os.getcwd())
          reload(botocore)
          reload(boto3)

For Lambda

As a prerequisite for Lambda, it is necessary to attach a zip file including the package file of the boto3 / botocore in Lambda Layer.
You can use the latest Boto3 by specifying /opt/python/botocore/data for AWS_DATA_PATH environment variable. The following python code defines environment variables by processing Lambda code, but it is also possible to set environment variables in Lambda settings.


          import os

          os.environ["AWS_DATA_PATH"] = "/opt/python/botocore/data"

In Lambda, Lambda Layer is mounted on /opt/python , and the module contained in the layer is prioritized.Unlike SSM Automation, there is no need to Reload Boto3 or Botocore.

Confirmation/verification result (for SSM Automation)

As described in the following class method, Botocore is set in $ {Botocore root directory}/data , ~/.aws/models and AWS_DATA_PATH When a client is generated in Boto3. I am reading the model file.
- Reference: 非最新のboto3で最新のAWSサービスを利用する方法 | DevelopersIO

As you can see from the directory where the Botocore is installed, the model file has a file configuration as follows.


          # For Control Tower
          botocore/data
          └── controltower
               └── 2018-05-10
                   ├── endpoint-rule-set-1.json.gz
                   ├── paginators-1.json # The list of APIs that can be used here is listed
                   └── service-2.json.gz

If this aws_data_path environment variable is set, you can use the latest model file because it is loaded with the highest priority.
- Reference: Configuration - Boto3 1.34.108 documentation z

However, even if this environment variable is set, Boto3 and Botocore are not the latest versions in SSM Automation.
In Lambda, the module contained in Lambda Layer is given preferentially, but in the case of SSM Automation, the module contained in the zip file, including the Attachment, is not loaded preferentially.
In the case of SSM Automation, the file including the Attachment is expanded to the /tmp/${random character string} .The following output is the JSON format in Sys.path, but the last priority of files in Attachment is the priority.


          {
              "sys path": [
                  "/var/task",
                  "/opt/python/lib/python3.8/site-packages",
                  "/opt/python",
                  "/var/runtime",
                  "/var/lang/lib/python38.zip",
                  "/var/lang/lib/python3.8",
                  "/var/lang/lib/python3.8/lib-dynload",
                  "/var/lang/lib/python3.8/site-packages",
                  "/opt/python/lib/python3.8/site-packages",
                  "/opt/python",
                  "/tmp/4aa790cf-cb0b-48b6-80f0-3ed7f82f228c-2024-03-17-07-15-31"
              ]
          }

Therefore, in the case of SSM Automation, it is necessary to add a current directory to Sys.path and Reload Boto3 and Botocore.


          import os
          import sys
          import boto3
          import botocore

          sys.path.insert(0, os.getcwd())
          reload(botocore)
          reload(boto3)

SSM Automation Terraform code used for verification

In order to perform the version upgrade of Boto3 of SSM Automation, I used the following file configuration Terraform code.


          ├── src
          │   ├── customer_script.py
          │   └── modules
          │       └── requirements.txt
          └── ssm.tf

src/customer_script.py

In SSM Automation, in addition to the process of upgrading Boto3, log output is also considered.
- You can use the ListHandler class to store the log output in the list and return it to the end to display the log on the SSM Automation execution screen.


          import logging
          import json
          import os
          import sys
          from importlib import reload

          import boto3
          import botocore

          AUTOMATION_NAME = "sample-automation"
          logger = logging.getLogger(__name__)


          class ListHandler(logging.Handler):
              def __init__(self):
                  super().__init__()
                  self.log_messages = []
                  log_format = json.dumps(
                      {
                          "logLevel": "%(levelname)s",
                          "time": "%(asctime)s",
                          "line": "%(filename)s: %(lineno)d",
                          "message": "%(message)s",
                      },
                      indent=2,
                      ensure_ascii=False,
                  )
                  self.formatter = logging.Formatter(log_format)

              def emit(self, record):
                  self.log_messages.append(json.loads(self.format(record)))


          def handler(event, context):
              list_handler = ListHandler()
              logger.addHandler(list_handler)
              logger.setLevel(logging.DEBUG)

              os.environ["AWS_DATA_PATH"] = f"{os.getcwd()}/botocore/data"
              session = boto3.Session()
              logger.debug({"botocore paths": session._loader.search_paths})

              sys.path.insert(0, os.getcwd())
              reload(botocore)
              reload(boto3)

              logger.debug({"sys path": sys.path})
              logger.debug({"current directory": os.getcwd()})
              logger.debug({"child directory": os.listdir()})
              logger.debug({"botocore path": os.path.dirname(botocore.__file__)})
              logger.debug({"boto3 path": boto3.__file__})
              logger.debug({"boto3 version": boto3.__version__})
              logger.debug({"botocore path": botocore.__file__})
              logger.debug({"botocore version": botocore.__version__})

              try:
                  client = session.client("controltower")
                  logger.info(client.list_baselines())

              except Exception as e:
                  logger.error(f"Failed to list_baselines: {e}")
                  raise Exception(list_handler.log_messages)

              # max length of ssm automation output is 100 KB(102400 characters)
              return list_handler.log_messages


          if __name__ == "__main__":
              handler(event={}, context={})

src/modules/requirements.txt

The version of the Boto3, Botocore used to verify the version of the SSM Automation Boto3 is described.
- It is not particularly version, but uses the latest version (1.34.64 at the time of writing).


          boto3
          botocore

ssm.tf

This is a Terraform code used to verify the version upgrade of Boto3 of SSM Automation.
- Comments such as #tfsec:ignore:aws-s3-encryption-customer-key , etc. are comments to disable TFSec verification.
- Using null_resource , dataarchive_file , and aws_s3_object , we are creating a ZIP file including the package file of Boto3 / Botocore in the Attachment of SSM Automation.
- When coding with CloudFormation, it does not support the aws cloudformation package command, so you need to manually create a ZIP file and get a SHA256 hash value.Terraform is convenient because it is flexible.



          data "aws_caller_identity" "current" {}

          #tfsec:ignore:aws-s3-encryption-customer-key
          #tfsec:ignore:aws-s3-enable-bucket-encryption
          #tfsec:ignore:aws-s3-enable-bucket-logging
          resource "aws_s3_bucket" "bucket" {
            bucket = "sample-automation-${data.aws_caller_identity.current.account_id}"
          }

          resource "aws_s3_bucket_versioning" "bucket_versioning" {
            bucket = aws_s3_bucket.bucket.id
            versioning_configuration {
              status = "Enabled"
            }
          }

          resource "aws_s3_bucket_public_access_block" "bucket_public_access_block" {
            bucket = aws_s3_bucket.bucket.id
            block_public_acls       = true
            block_public_policy     = true
            ignore_public_acls      = true
            restrict_public_buckets = true
          }

          resource "null_resource" "install_requirements" {
            triggers = {
              always_run = "${timestamp()}"
            }

            provisioner "local-exec" {
              command = "pip install -r ${path.root}/src/modules/requirements.txt -t ${path.root}/src/modules"
            }
          }

          data "archive_file" "file" {
            depends_on = [ null_resource.install_requirements ]
            type        = "zip"
            source_dir  = "${path.root}/src/modules"
            output_path = "${path.root}/src/modules.zip"
          }

          resource "aws_s3_object" "object" {
            bucket = aws_s3_bucket.bucket.id
            key    = "modules.zip"
            source = data.archive_file.file.output_path
          }

          resource "aws_ssm_document" "runbook" {
            depends_on = [ aws_s3_object.object ]

            name            = "sample-automation"
            document_type   = "Automation"
            document_format = "JSON"
            attachments_source {
              key = "S3FileUrl"
              name = "modules.zip"
              values = [
                "https://${aws_s3_bucket.bucket.bucket}.s3.amazonaws.com/modules.zip"
              ]
            }

            content = jsonencode({
              assumeRole      = aws_iam_role.role.arn
              schemaVersion = "0.3"
              description   = "TestBoto3VersionUp"
              mainSteps = [
                {
                  name   = "TestBoto3VersionUp"
                  action = "aws:executeScript"
                  inputs = {
                    Runtime = "python3.8"
                    Handler = "handler"
                    Script  = "${file("${path.root}/src/customer_script.py")}"
                    Attachment = "modules.zip"
                  }
                }
              ]
              files = {
                "modules.zip" = {
                  checksums = {
                    sha256 = filesha256(data.archive_file.file.output_path)
                  }
                }
              }
            })
          }

          resource "aws_iam_role" "role" {
            name = "sample-automation-role"
            assume_role_policy = jsonencode({
              Version = "2012-10-17",
              Statement = [
                {
                  Action = "sts:AssumeRole",
                  Effect = "Allow",
                  Principal = {
                    Service = "ssm.amazonaws.com"
                  }
                }
              ]
            })
            managed_policy_arns = [
              "arn:aws:iam::aws:policy/service-role/AmazonSSMAutomationRole",
              "arn:aws:iam::aws:policy/AdministratorAccess"
            ]
          }

Terraform code used to verify the version upgrade of Boto3 of Lambda

I used the Terraform code with the following file configuration to perform Lambda’s Boto3 version upgrade verification.


          ├── src
          │   ├── function
          │   │   └── main.py
          │   └── layer
          │   │   ├── python # Store the package file of Boto3 / Botocore by processing Terraform under this directory
          │   │   └── requirements.txt
          └── lambda.tf

src/function/main.py

The Python code of Lambda is simple.


          import logging
          import os
          import sys

          import boto3
          import botocore

          logger = logging.getLogger(__name__)


          def handler(event, context):
              logger.setLevel(logging.DEBUG)

              os.environ["AWS_DATA_PATH"] = "/opt/python/botocore/data"
              session = boto3.Session()
              logger.debug({"botocore paths": session._loader.search_paths})

              logger.debug({"sys path": sys.path})
              logger.debug({"botocore path": os.path.dirname(botocore.__file__)})
              logger.debug({"boto3 path": boto3.__file__})
              logger.debug({"boto3 version": boto3.__version__})
              logger.debug({"botocore path": botocore.__file__})
              logger.debug({"botocore version": botocore.__version__})

              client = session.client("controltower")
              logger.info(client.list_baselines())


          if __name__ == "__main__":
              handler(event={}, context={})

src/layer/requirements.txt

The package file to be included in Lambda Layer is described.
- It is not particularly version, but uses the latest version (1.34.64 at the time of writing).


          boto3
          botocore

lambda.tf

This is a Terraform code used to verify the version upgrade of Boto3 of Lambda.
- Like SSM Automation, we use null_resource and data.archive_file to create a ZIP file including Boto3 / Botocore package files in Lambda Layer.



          data "archive_file" "function" {
            depends_on  = [null_resource.install_requirements]
            type        = "zip"
            source_dir  = "${path.root}/src/function"
            output_path = "${path.root}/src/function.zip"
          }

          resource "aws_lambda_function" "function" {
            function_name    = "sample-function"
            filename         = data.archive_file.function.output_path
            source_code_hash = data.archive_file.function.output_base64sha256
            handler          = "main.handler"
            runtime          = "python3.11"
            role             = aws_iam_role.role.arn
            # os.environ["AWS_DATA_PATH"] is OK even if you set the following environment variables
            # environment {
            #   variables = {
            #     AWS_DATA_PATH = "/opt/python/botocore/data"
            #   }
            # }
            layers        = [aws_lambda_layer_version.layer_version.arn]
          }

          resource "null_resource" "install_requirements" {
            triggers = {
              always_run = "${timestamp()}"
            }

            provisioner "local-exec" {
              command = "pip install -r ${path.root}/src/layer/requirements.txt -t ${path.root}/src/layer/python"
            }
          }

          data "archive_file" "layer" {
            depends_on  = [null_resource.install_requirements]

            type        = "zip"
            source_dir  = "${path.root}/src/layer"
            output_path = "${path.root}/src/layer.zip"
          }

          resource "aws_lambda_layer_version" "layer_version" {
            layer_name = "sample-layer"
            filename  = data.archive_file.layer.output_path
            compatible_runtimes = ["python3.11"]
            source_code_hash = data.archive_file.layer.output_base64sha256
          }

          resource "aws_iam_role" "role" {
            name = "sample-lambda-role"
            assume_role_policy = jsonencode({
              Version = "2012-10-17",
              Statement = [
                {
                  Action = "sts:AssumeRole",
                  Effect = "Allow",
                  Principal = {
                    Service = "lambda.amazonaws.com"
                  }
                }
              ]
            })
            managed_policy_arns = [
              "arn:aws:iam::aws:policy/AdministratorAccess"
            ]
          }