LambdaforAutoScalingControlOnCodePipelineWithCrossAccount

from botocore.exceptions import ClientError

import json
import boto3
import traceback
import os

client = boto3.client('autoscaling')

def get_cross_account_connection():
    sts_connection = boto3.client('sts')
    try:
        account_pipeline = sts_connection.assume_role(
            RoleArn="arn:aws:iam::{}:role/{}".format(os.environ['CROSS_ACCOUNT_ID'], os.environ['CROSS_ACCOUNT_ASSUME_ROLE']),
            RoleSessionName="account_pipeline"
        )
    except ClientError as e:
        print(e.response['Error']['Message'])
        raise e
    
    # create service client using the assumed role credentials, e.g. S3
    try:
        code_pipeline = boto3.client(
            'codepipeline',
            aws_access_key_id=account_pipeline['Credentials']['AccessKeyId'],
            aws_secret_access_key=account_pipeline['Credentials']['SecretAccessKey'],
            aws_session_token=account_pipeline['Credentials']['SessionToken'],
        )
        return code_pipeline
    except ClientError as e:
        print(e.response['Error']['Message'])
        raise e
        
def resume_processes(auto_scaling_group_name):
    """
    Suspend processes of AutoScalingGroup
    Args:
        auto_scaling_group_name: auto scaling group name
    """
    # Processes
    processes=['ReplaceUnhealthy','AZRebalance','AlarmNotification','ScheduledActions']
    
    print("INFO: Suspend target: {}".format(auto_scaling_group_name))
    try:
        response = client.resume_processes(
            AutoScalingGroupName=auto_scaling_group_name,
            ScalingProcesses=processes
        )
    except ClientError as e:
        print(e.response['Error']['Message'])
        raise e

def put_job_success(job, message):
    """
    Notify CodePipeline of a successful job
    Args:
        job: The CodePipeline job ID
        message: A message to be logged relating to the job status
    Raises:
        Exception: Any exception thrown by .put_job_success_result()
    """
    print(message)
    get_cross_account_connection().put_job_success_result(jobId=job)

def put_job_failure(job, message):
    """
    Notify CodePipeline of a failed job
    Args:
        job: The CodePipeline job ID
        message: A message to be logged relating to the job status
    Raises:
        Exception: Any exception thrown by .put_job_failure_result()
    """
    print(message)
    get_cross_account_connection().put_job_failure_result(jobId=job, failureDetails={'message': message, 'type': 'JobFailed'})

def get_user_params(job_data):
    """
    Decodes the JSON user parameters and validates the required properties.
    Args:
        job_data: The job data structure containing the UserParameters string which should be a valid JSON structure
    Returns:
        The JSON parameters decoded as a dictionary.
    Raises:
        Exception: The JSON can't be decoded or a property is missing.
    """
    try:
        # Get the user parameters which contain the stack, artifact and file settings
        user_parameters = job_data['actionConfiguration']['configuration']['UserParameters']
        decoded_parameters = json.loads(user_parameters)

    except Exception as e:
        # We're expecting the user parameters to be encoded as JSON
        # so we can pass multiple values. If the JSON can't be decoded
        # then fail the job with a helpful message.
        raise Exception('UserParameters could not be decoded as JSON')

    if 'auto_scaling_group_name' not in decoded_parameters:
        # Validate that the auto scaling group name is provided, otherwise fail the job
        # with a helpful message.
        raise Exception('Your UserParameters JSON must include a auto scaling group name')

    return decoded_parameters

def lambda_handler(event, context):
    """
    The Lambda function handler

    Args:
        event: The event passed by Lambda
        context: The context passed by Lambda
    """
    job_id = event['CodePipeline.job']['id']            # Extract the Job ID
    job_data = event['CodePipeline.job']['data']        # Extract the Job Data
    name = ""
    try:
        params = get_user_params(job_data)              # Extract the params
        name = params['auto_scaling_group_name']        # Operation target 
        resume_processes(name)
    except Exception as e:
        # If any other exceptions which we didn't expect are raised
        # then fail the job and log the exception message.
        print('Function failed due to exception.')
        print(e)
        traceback.print_exc()
        put_job_failure(job_id, 'Function exception: ' + str(e))
        return False

    # Update pipeline state
    put_job_success(job_id, "Suspend autoscaling is Success. target: {}".format(name))
    return True