StepFunctions and State Machine

List state machines

$ aws stepfunctions list-state-machines --profile nonprod_admin | jq .[][].name | grep -i 'dev\|test\|nonprod'
"dev-approval-workflow"
"dev-get-accounts"
"ami-lifecycle-dev-ami-statemachine"

Core Step Function

Code for the Lifecycle checking step function

{
    "Comment": "StateMachine Framework to get lifecycle code  working in states.",
    "StartAt": "StartState",
    "States": {
      "StartState":{
            "Comment": "Passes input data",
            "Type": "Pass",
            "Result": {
                      "AMIage": "30"
                      },
             "ResultPath": "$",
             "Next": "AMIolderThanXdays"
      },
      "AMIolderThanXdays": {
            "Comment": "Local State comment",
            "Type": "Task",
            "Resource": "arn:aws:lambda:eu-west-1:05772692xxxx:function:AJS-test-statemachine-AMIolderThanX",
            "ResultPath": "$.AMIolderThanX",
            "Next": "AmisStillToProcess_Choice"
        },
        "AmisStillToProcess_Choice": {
            "Type": "Choice",
            "Choices": [
                {
                    "Variable": "$.AMIolderThanX.continue",
                    "BooleanEquals": true,
                    "Next": "AMIgetAccounts"
                }
            ],
            "Default": "Done"
        },
        "AMIgetAccounts": {
            "Type": "Task",
            "Resource": "arn:aws:lambda:eu-west-1:05772692xxxx:function:AJS-testStateMachine-getAccounts",
            "InputPath": "$.AMIolderThanX",
            "ResultPath": "$.Accounts",
            "Next": "AccountsStillToProcess_Choice"
        },
      "AccountsStillToProcess_Choice": {
            "Type": "Choice",
            "Choices": [
                {
                    "Variable": "$.Accounts.continueAcc",
                    "BooleanEquals": true,
                    "Next": "AMIgetUsage"
                }
            ],
            "Default": "NextAmi"
        },
        "AMIgetUsage": {
            "Type": "Task",
            "Resource": "arn:aws:lambda:eu-west-1:05772692xxxx:function:AJS-test-StateMachine-AMIgetUsage",
            "InputPath": "$.Accounts",
            "ResultPath": "$.Accounts",
            "Next": "AMIusageAddToNotifier"
        },
        "AMIusageAddToNotifier":{
            "Type": "Task",
            "Resource": "arn:aws:lambda:eu-west-1:05772692xxxx:function:AJS-test-StateMachine-AMI_Account_to_Notifier",
            "InputPath": "$.Accounts",
            "ResultPath": "$.Accounts",
            "Next": "AccountsStillToProcess_Choice"
        },
 
 
 
        "NextAmi": {
            "Type": "Task",
            "Resource": "arn:aws:lambda:eu-west-1:05772692xxxx:function:AJS-test-statemachine-iterator",
            "InputPath": "$.AMIolderThanX",
            "ResultPath": "$.AMIolderThanX",
            "Next": "AmisStillToProcess_Choice"
        },
        "Done": {
            "Type": "Pass",
            "End": true
        }
    }
}

IAM policy

AndrewLambdaTestPolicy

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "VisualEditor0",
            "Effect": "Allow",
            "Action": [
                "states:ListStateMachines",
                "ec2:DescribeInstances",
                "states:ListActivities",
                "states:CreateActivity",
                "autoscaling:DescribeLaunchConfigurations",
                "states:StopExecution",
                "ec2:DescribeImages",
                "states:SendTaskSuccess",
                "states:SendTaskFailure",
                "autoscaling:DescribeAutoScalingGroups",
                "ec2:DescribeImageAttribute",
                "states:SendTaskHeartbeat",
                "states:CreateStateMachine",
                "sts:AssumeRole"
            ],
            "Resource": "*"
        },
        {
            "Sid": "VisualEditor1",
            "Effect": "Allow",
            "Action": "states:*",
            "Resource": [
                "arn:aws:states:*:*:activity:*",
                "arn:aws:states:*:*:execution:*:*",
                "arn:aws:states:*:*:stateMachine:*"
            ]
        }
    ]
}

AmazonSQSFullAccess

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": [
                "sqs:*"
            ],
            "Effect": "Allow",
            "Resource": "*"
        }
    ]
}

AWSLambdaBasicExecutionRole

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "logs:CreateLogGroup",
                "logs:CreateLogStream",
                "logs:PutLogEvents"
            ],
            "Resource": "*"
        }
    ]
}

AWSLambdaRole

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "lambda:InvokeFunction"
            ],
            "Resource": [
                "*"
            ]
        }
    ]
}

Lambda Functions

AJS-test-statemachine-AMIolderThanX

import os
import boto3
import datetime
from datetime import date
from datetime import datetime
 
 
# Simple Lambda to return ami list to emulate ami older than X days
 
def checkAMIinUse(AMIage, context):
    region = os.environ['REGION']
    ec2 = boto3.client('ec2',region)
 
    AGEthreshold = int(os.environ['AGEthreshold'])
    Regex = os.environ['REGEX']
 
    AMIlist = []
    # print (AGEthreshold, Regex)
 
    # get list of all AMI owned by this account, matching the regex
    AMIResponse = ec2.describe_images(Filters=[{'Name': 'name', 'Values': [Regex]}, ], Owners=['self'])
    for i in AMIResponse['Images']:
 
        AMIname = i['Name']
        AMIimageID = i['ImageId']
        AMIcreationDate = i['CreationDate']
 
        # convert unicode string to timedate object
        AMICreationDateTime = datetime.strptime((i['CreationDate']), '%Y-%m-%dT%H:%M:%S.%fZ')
 
        # only needs days not time, so use date method, likewise now is just days
        AMICreationDate = AMICreationDateTime.date()
        now = date.today()
 
        # Get age of ami and convert timedelta to contain only days
        AMIage = (now - AMICreationDate).days
 
        if AMIage >= AGEthreshold:
            olderThanThreshold = True
            AMIlist.append(AMIimageID)
        else:
            olderThanThreshold = False
 
        # print ("AMI", AMIimageID, "Age Threshold ", AGEthreshold, "AMI age ",AMIage, "olderthanthres ", olderThanThreshold)
        # print (AMIlist)
 
    # resultDict = {"OlderThanThreshold":olderThanThreshold, "AMIage":AMIage}
    # return (AMIlist)
 
 
    index = 0
    count = len(AMIlist)
 
    result = {
        "ami_ids": AMIlist,
        "index": index,
        "count": count,
        "continue": index < count
    }
 
    return (result)

AJS-testStateMachine-getAccounts

# Simple Lambda to return account numbers for ami shared with.
# AccoutID  == UserID
import boto3
import sys
 
 
def GetAMIaccounts(AMIinfo, context):
    ec2 = boto3.client('ec2')
 
    AMIlist =  (AMIinfo['ami_ids'])
    Index = int(AMIinfo['index'])
 
    image =  AMIlist[Index]
 
    # aws ec2 describe-image-attribute --image-id ami-07cf57ebf50e78466 --attribute launchPermission --profile nonprod_admin
    #try:
    response = ec2.describe_image_attribute(ImageId=image, Attribute='launchPermission')
 
    Accounts =  response['LaunchPermissions']
 
    accountlist = []
    for account in Accounts:
        accountlist.append(account['UserId'])
 
    # print (AMIimageID['ImageID'],  accountlist)
 
    indexAcc = 0
    countAcc = len(accountlist)
    if countAcc == 0:
        continueAcc = False
    else:
        continueAcc = True
 
    response = {}
    response['ami_id'] = image
    response['account_id'] = accountlist
    response['indexAcc'] = indexAcc
    response['countAcc'] = countAcc
    response['continueAcc'] = continueAcc
    print (response)
 
 
    return (response)

AJS-test-statemachine-iterator

def lambda_handler(event, context):
    index = event['index']
    count = event['count']
 
    index += 1
    event['index'] = index
    # "index < count" evalutes to either true or false
    # to control the step loop.
    event['continue'] = index < count
 
    return event

AJS-test-StateMachine-AMIgetUsage

import os
import boto3
 
def lambda_handler(event, context):
    sts_client = boto3.client('sts')
    print (event)
 
    index = int(event['indexAcc'])
    count = int(event['countAcc'])
    ami_id = event['ami_id']
    accountToAssume = event['account_id'][index]
    roleToAssume = os.environ['roleToAssume']
    # print (index, type(index))
    # print ("Account is ",accountToAssume, "roleToAssume", roleToAssume, "AMI id", ami_id)
 
    assumed_role_object=sts_client.assume_role(
        RoleArn="arn:aws:iam::"+accountToAssume+":role/"+roleToAssume,
        RoleSessionName="AMI_Lifecycle_StateMachine"
    )
 
    # print ("RoleArnTest", RoleArnTest )
 
    credentials=assumed_role_object['Credentials']
    # print (credentials)
 
 
    ec2client = boto3.client(
        'ec2',
        aws_access_key_id=credentials['AccessKeyId'],
        aws_secret_access_key=credentials['SecretAccessKey'],
        aws_session_token=credentials['SessionToken']
        )
 
    autoscaling_client = boto3.client(
        'autoscaling',
        aws_access_key_id=credentials['AccessKeyId'],
        aws_secret_access_key=credentials['SecretAccessKey'],
        aws_session_token=credentials['SessionToken']
        )
 
    # Start of real code, we have already selected the correct account from the list of accounts to switch to,
    # but there is only one AMI being passed in.
 
    # aws ec2 describe-instances --filters "Name=image-id, Values=ami-0e12cbde3e77cbb98" --query 'Reservations[*].Instances[*].[InstanceId]'  --profile nonprod_admin
    #EC2InUseResponse = ec2client.describe_instances(Filters=[{'Name': 'image-id', 'Values': [ami_id['ImageID']]}])
    EC2InUseResponse = ec2client.describe_instances(Filters=[{'Name': 'image-id', 'Values': [ami_id]}])
 
    # list comprehension, returns list of instances
    EC2instance_ids = [
        i["InstanceId"]
        for r in EC2InUseResponse["Reservations"]
        for i in r["Instances"]
    ]
    # print ('checkAMIinUse: Instances built from ',ami_id, ' are ', EC2instance_ids)
    if len(EC2instance_ids) == 0:
        EC2AMIinUse = False
    else:
        EC2AMIinUse = True
        # print ('checkAMIinUse: Instances ', EC2AMIinUse)
 
    # End of ec2 instance detection
    # print ( '=' * 10 )
 
    # get all ASG and for ech, get the Launch config.
    # check each LC for the AMI in AMIimageID
    # this gets all asg but have to match lc passed to function within it.
    # aws autoscaling describe-auto-scaling-groups --profile nonprod_admin --auto-scaling-group-names "AJS asg1"
 
    # print ('checkAMIinUse: ASG Check if AMI is in use by ASG')
    ASGresponse = autoscaling_client.describe_auto_scaling_groups()
 
    ASGAMIinUse = False
    ASGlist = []
    for i in ASGresponse['AutoScalingGroups']:
        # get launch config from ASG
        LC_Name = i['LaunchConfigurationName']
        ASG_Name = i['AutoScalingGroupName']
        # print ('LC_Name is:-', LC_Name, 'ASG_Name is:- ', ASG_Name)
 
        # get all launch conf info and select ImageID 
        LCResponse = autoscaling_client.describe_launch_configurations(LaunchConfigurationNames=[LC_Name, ], )
        #LCResponse = getAllLaunchConf(LC_Name)
        LC_object = LCResponse['LaunchConfigurations']
 
        for item in LC_object:
            # Check if imagesIDs are the same , if they are, add LaunchConf to list
            # print ('Testing for AMIimageID ', AMIimageID, ' in ', item['ImageId'])
            #if AMIimageID['ImageID'] == item['ImageId']:
            if ami_id == item['ImageId']:
                # print('True')
                ASGlist.append(item['LaunchConfigurationName'])
                ASGAMIinUse = True
 
                # This gets the EC2 instances fired up as part of autoscaling
                # $ aws autoscaling describe-auto-scaling-groups --profile nonprod_admin --auto-scaling-group-name 'AJS-asg1' --query 'AutoScalingGroups[*].Instances[*].InstanceId'
 
                # Create an empty list and populate with all the instances built as part of an asg
                ASGinstances = []
                instanceInASG = autoscaling_client.describe_auto_scaling_groups(AutoScalingGroupNames=[ASG_Name])
 
                for i in instanceInASG['AutoScalingGroups']:
                    for inst in (i['Instances']):
                        # print (inst['InstanceId'])
                        ASGinstances.append(inst['InstanceId'])
 
 
                #print ('All EC2 instances ', EC2instance_ids)
                #print ('ASGinstances:- ', ASGinstances)
                for instASG in ASGinstances:
                    # print (instASG)
                    # remove asg drived instance from list of EC2 instances
                    EC2instance_ids.remove(instASG)
 
                break
            break
 
    # print ('checkAMIinUse: ASG  Ami InUse:- ', ASGAMIinUse, ASGlist  )
 
    if EC2AMIinUse or ASGAMIinUse is True:
        AMIinUse = True
    else:
        AMIinUse = False
    # print ( '+' * 10 )
 
 
    # Finish off, inc index and write info back to state machine
    index += 1
    event['indexAcc'] = index
    # "index < count" evalutes to either true or false
    # to control the step loop.
    event['continueAcc'] = index < count
    event['AMIusage'] = {'accountID':[accountToAssume], 'InstanceID':EC2instance_ids, 'ASG_ID':ASGlist}
 
 
    return (event)
    # return (AMIinUse, EC2instance_ids, ASGlist)
 
"""
{
  "ami_id": "ami-07cf57ebf50e78466",
  "account_id": [
    "532982424333",
    "057726927330",
    "1234567890"
  ],
  "indexAcc": 1,
  "countAcc": 3,
  "continueAcc": true,
  "AMIusage": {
    "accountID": [
      "532982424333"
    ],
    "InstanceID": [
      "inst123456",
      "inst654321"
    ],
    "ASG_ID": [
      "asg098765",
      "asg456789"
    ]
  }
}
"""
Env variables
roleToAssume AMI-lifecycle-usage

Role needs autoscaling and ec2 access in each account to assume

AJS-test-StateMachine-AMI_Account_to_Notifier

import os
import boto3
 
def lambda_handler(event, context):
    sqs = boto3.client('sqs')
    queue_url = os.environ['SQS_Queue']
    MessageBody = os.environ['MessageBody']
 
    index = event['indexAcc']
    count = event['countAcc']
    ami = event['ami_id']
 
    """
    "AMIusage": {
     "accountID": [
      "532982424333"
     ],
     "InstanceID": [
       "inst123456",
       "inst654321"
     ],
     "ASG_ID": [
       "asg098765",
       "asg456789"
     ]
    } """
    usage = event['AMIusage']
    usage.update({'AMI_ID':ami})
 
    account = usage['accountID']
    instances = usage['InstanceID']
    autoScaling = usage['ASG_ID']
    AMIage = '30'
 
    # BUG!! don't send message if instances AND asg are both empty.
    print (instances, len(instances), autoScaling, len(autoScaling))
    if len(instances) == 0 and len(autoScaling) == 0:
        print ('Quitting early')
        return (event)
 
    # Send message to SQS queue
    response = sqs.send_message(
        QueueUrl=queue_url,
        DelaySeconds=0,
        MessageAttributes={
            'AMIage': {
                'DataType': 'Number',
                'StringValue': AMIage
            },
            'AMIimageID': {
                'DataType': 'String',
                'StringValue': ami
            },
            'accountID': {
                'DataType': 'String',
                'StringValue': str(account)
            },
            'instances': {
                'DataType': 'String',
                'StringValue': str(instances)
            },
            'autoScaling': {
                'DataType': 'String',
                'StringValue': str(autoScaling)
            }
 
            },
        MessageBody = MessageBody
    )
 
    return (event)
Env variables
MessageBody Your account references an AMI older than the threshold in either EC2 instances or AutoScaling Launch Configurations. Please see Message Attributes for details.
SQS_Queue https://sqs.eu-west-1.amazonaws.com/057726927330/AJS-testQueue
 
aws/statemachine.txt · Last modified: 27/01/2022 15:57 by andrew