Lambda S3->AMI converter

bootstrap.sh

#!/bin/bash
# Copyright 2021 Josh Snyder
# See license and documentation in lambda_function.py
set -x -o errexit -o nounset

ebs_name() { 
  nvme id-ctrl -o binary $1  | cut -c3073-3104 | tr -d ' '
}

export AWS_DEFAULT_REGION=$AWS_REGION  

# Give devices pretty names
for dev in /dev/nvme*n*; do
  source=/dev/$(ebs_name $dev)
  test -e $source || ln -s $dev $source
done

# Make S3 fast
mkdir /root/.aws
cat<<EOF > /root.aws
[default]
s3 =
    multipart_chunksize = 128MB
    max_concurrent_requests = 4
EOF

aws s3 cp --no-progress s3://<my-artifacts-bucket>/test1.lz4 - | lz4 -f -d - /dev/sdh

volume_id=vol-$(nvme id-ctrl -o binary /dev/sdh | cut -c 8-24)
snapshot_id=$(
  aws ec2 create-snapshot --volume-id "${volume_id}" |
  jq -r ".SnapshotId"
)

aws ec2 wait snapshot-completed --snapshot-id "$SNAPSHOT_ID"

ami_name=$(uuidgen)
image_id=$(
  aws ec2 register-image --name "${ami_name}" --root-device-name /dev/sda1 \
    --boot-mode uefi \
    --block-device-mapping "DeviceName=/dev/sda1,Ebs={SnapshotId=$snapshot}" \
    --ena-support --virtualization-type hvm --architecture x86_64 |
  jq -r ".ImageId"
)

echo $image_id

lambda_function.py

"""Lambda function that converts S3 blobs into AMIs.
==========================
Copyright 2021 Josh Snyder

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==========================

This lambda function essentially reimplements `aws ec2 import-image`, but
better in every way.

To use:
 1) Make an S3 bucket for the VM artifacts, copy it into bootstrap.sh (line 25)
 2) Make an IAM instance profile for the builder instance; modify
    IamInstanceProfile below. Must be able to access the S3 bucket and register
    AMIs.
 3) Make a VPC security group: no ingress, egress allowed to S3 and the AWS EC2
    API. Put it at SecurityGroupIds below.
 4) If you want the ability to debug, copy your favorite SSH public key into
    ssh_authorized_keys below.
 5) Upload an object to s3://<my-artifacts-bucket>/test.lz4
 6) Trigger the lambda
"""
import os
import boto3
import json
import sys
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from pathlib import Path

region = os.environ["AWS_REGION"]

autoscaling = boto3.client("autoscaling", region_name=region)
ec2 = boto3.client('ec2', region_name=region)
s3 = boto3.client('s3')

USERDATA = """#cloud-config

ssh_authorized_keys:
 - ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIKr4DFWVEoLCTgjtzl3wT+JnYnDojJAS/4hsFww4n/R8
 
bootcmd:
 - dd if=/dev/zero count=1 of=$(lsblk -n -o pkname /dev/disk/by-label/cloudimg-rootfs)
 - systemctl mask snapd
 - touch /var/cache/pollinate/seeded

power_state:
  delay: +1
  mode: poweroff
  message: cloud-init shutting down
  timeout: 300
"""

BOOTSTRAP_PATH = Path(os.environ['LAMBDA_TASK_ROOT'] + "/bootstrap.sh")
BOOTSTRAP_SCRIPT = BOOTSTRAP_PATH.read_text()

def compose_user_data():
    ret = MIMEMultipart()
    ret.attach(
        MIMEText(USERDATA, "cloud-config", sys.getdefaultencoding())
    )
    ret.attach(
        MIMEText(BOOTSTRAP_SCRIPT, "x-shellscript", sys.getdefaultencoding())
    )
    return str(ret)


TASK_TAG="bake-ami"


RUN_INSTANCES_PARAMS = dict(
    ImageId="ami-00482f016b2410dc8", # Ubuntu 21.10 hvm-ssd
    InstanceType="t3.nano",
    MinCount=1,
    MaxCount=1,
    TagSpecifications=[
        dict(ResourceType="instance", Tags=[dict(Key="task", Value=TASK_TAG)]),
        dict(ResourceType="volume", Tags=[dict(Key="Owner", Value=TASK_TAG)]),
    ],
    UserData=compose_user_data(),
    IamInstanceProfile=dict(Name="bake-ami.ec2"),
    InstanceInitiatedShutdownBehavior="terminate",
    SecurityGroupIds=["sg-0257605ee80f7fa51"],

)


def lambda_handler(event, context):
    run_params = dict(RUN_INSTANCES_PARAMS)
    
    run_params["BlockDeviceMappings"] = [
        dict(
            DeviceName="/dev/sdh",
            Ebs=dict(
                DeleteOnTermination=True,
                VolumeSize=1    ,
                VolumeType="gp2"
            ),
        ),
    ]
    
    resp = ec2.run_instances(**run_params)
    
    body = dict()
    for ins in resp["Instances"]:
        body["id"] = ins["InstanceId"]
        for interface in ins["NetworkInterfaces"]:
            for addr in interface["Ipv6Addresses"]:
                body["address"] = addr["Ipv6Address"]
                
    return dict(
        statusCode=200,
        headers={"Content-Type": "application/json"},
        body=json.dumps(run_params["UserData"]),
    )