#!/usr/bin/env python3
import boto3

ec2 = boto3.client('ec2', region_name='us-west-2')
sts_client = boto3.client('sts')

# Call the assume_role method of the STSConnection object and pass the role
# ARN and a role session name.
assumed_role_object = sts_client.assume_role(
    RoleArn="arn:aws:iam::369874303498:role/cl/app/crossaccount/amimanager/jenkins_role_policy",
    RoleSessionName="AssumeRoleSession1"
)
credentials = assumed_role_object['Credentials']
autoscaling = boto3.client('autoscaling',
                           aws_access_key_id=credentials['AccessKeyId'],
                           aws_secret_access_key=credentials['SecretAccessKey'],
                           aws_session_token=credentials['SessionToken'],
                           region_name='us-west-2')
ec2 = boto3.client('ec2',
                   aws_access_key_id=credentials['AccessKeyId'],
                   aws_secret_access_key=credentials['SecretAccessKey'],
                   aws_session_token=credentials['SessionToken'],
                   region_name='us-west-2')


def get_asg_instances(asgs):
    asg_instances = []
    for asg in asgs['AutoScalingGroups']:
        for instance in asg['Instances']:
            asg_instances.append(instance['InstanceId'])
    return asg_instances


def get_asg_instance_info(asgi, option=False):
    result = []
    for reservation in asgi.get('Reservations'):
        for instance in reservation.get('Instances'):
            if (option == False):
                result.append({'PrivateDnsName': instance['PrivateDnsName'],
                               'PrivateIpAddress': instance['PrivateIpAddress']})
            else:
                if (instance.get(option) is not None):
                    result.append({option: instance.get(option)})
                else:
                    print("Debug info: Instance doesn't have an PrivateIpAddress")
    return result
# For export as a function


def get_private_info(endpoint, option):
    asgs = autoscaling.describe_auto_scaling_groups(
        AutoScalingGroupNames=[endpoint])
    asgi_ids = get_asg_instances(asgs)
    asgi = ec2.describe_instances(InstanceIds=asgi_ids)
    test_data = get_asg_instance_info(asgi, option)
    return test_data


# As a separate script
if __name__ == "__main__":
    print('main')
    # asgs = autoscaling.describe_auto_scaling_groups(
    #     AutoScalingGroupNames=[AUTO_SCALING_GROUP])
    # asgi_ids = get_asg_instances(asgs)
    # asgi = ec2.describe_instances(InstanceIds=asgi_ids)
    # test_data = get_asg_instance_info(asgi)