elbtree

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import sys
import json
import argparse

try:
    import boto3
except:
    print "You need boto3. Please install boto3 like: 'pip install boto3'"
    sys.exit()

'''
ELBTree - ELBとその配下のインスタンスをツリー図ライクに表示

20170920 watanabe.seigo@classmethod.jp
'''

'''
データ構造:

dict:
    LoadBalancerName: LB名
    TargetGroupName: ターゲットグループ名(ALB/NLB)
    Type: classic/application/network

    Instances:
        [
            {
                Id: インスタンスID
                Name: インスタンス名(Nameタグ)
                Status: InService/OutOfService/Healthy/Unhealthy/unused
                AvailabilityZone: アベイラビリティゾーン
            }
        ]
'''

#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#
# ツリー表示時のルート
DESCRIPTION = {
    'classic'     : 'CLB (classic load balancer)',
    'application' : 'ALB (application load balancer)',
    'network'     : 'NLB (network load balancer)',
}

#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#
# インスタンスステータスのマーク表示
MARK = {
    "InService":   '--', "OutOfService": '//',
    "healthy":     '--', "unhealthy":    '//', "unused": '//',
    "__default__": '??',
}

#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#
class InstanceSummary():
    ''' インスタンス情報サマリ '''
    def __init__(self, instance_id, ec2s):
        '''
        self.instance_id ....... インスタンスID
        self.name .............. Nameタグ
        self.availabilityzone .. AvailabilityZone
        '''
        self.instance_id = instance_id

        # AZ, Nameタグ
        for ec2s_r in ec2s:
            instance_temp = [x for x in  ec2s_r['Instances'] if x['InstanceId'] == self.instance_id]
            if len(instance_temp) != 0:
                self.availabilityzone = instance_temp[0]['Placement']['AvailabilityZone']

                # Nameタグの情報取得
                tag_name = [x for x in instance_temp[0]['Tags'] if x['Key'] == 'Name']
                if len(tag_name) == 1 and tag_name[0]['Value'] != '':
                    self.name = tag_name[0]['Value']
                else:
                    self.name = None

                # 1回見つかればあとは不要
                break

#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#
def target_health_description(status_objct):
    ''' 'N/A' 以外のDescripttionを持っていたらそれを出力 '''
    if status_objct.has_key('Description') and status_objct['Description'] != 'N/A':
        return '{}'.format(status_objct['Description'])
    else:
        # 上記以外は None
        return None

#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#
def instance_infos(instance_id, instance_health_object, ec2s):
    ''' インスタンスの必要な情報を整理して出力 '''
    instance_summary_object = InstanceSummary(instance_id, ec2s)
    return {'Id':               instance_summary_object.instance_id,
            'Name':             instance_summary_object.name,
            'Status':           instance_health_object['State'],
            'Status_reason':    target_health_description(instance_health_object),
            'AvailabilityZone': instance_summary_object.availabilityzone
           }

#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#
def clb_data(aws, elb_datas, elb_list, elbs_all, ec2s):
    ''' CLB '''
    # ループ
    for elb_name in elb_list:
        # 対象ELBの情報抽出
        elb = [x for x in elbs_all if x['LoadBalancerName'] == elb_name][0]
        elb_data = {
            'LoadBalancerName': elb_name,
            'TargetGroupName': None,
            'Type': 'classic',
            'Instances': []
        }

        # インスタンスヘルス情報を取得
        instance_health = aws.client('elb').describe_instance_health(
            LoadBalancerName=elb_name
            )['InstanceStates']

        # 配下のインスタンス
        for instance in elb['Instances']:
            instance_id = instance['InstanceId']
            elb_data['Instances'] += [
                instance_infos(
                    instance_id,
                    [x for x in instance_health if x['InstanceId'] == instance_id][0],
                    ec2s)
                ]

        # 収集
        elb_datas += [elb_data]

    return elb_datas

#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#
def elbv2_data(aws, elb_datas, elb_list, elb_lbs_all, elb_tgs_all, ec2s):
    ''' ALB/NLB '''
    # ループ
    for elb_name in elb_list:
        # 対象ELBの情報抽出
        elb_lb = [x for x in elb_lbs_all if x['LoadBalancerName'] == elb_name][0]
        elb_tg = [x for x in elb_tgs_all if elb_lb['LoadBalancerArn'] in x['LoadBalancerArns']][0]
        elb_data = {
            'LoadBalancerName': elb_name,
            'TargetGroupName': elb_tg['TargetGroupName'],
            'Type': elb_lb['Type'],
            'Instances': []
        }

        # インスタンスヘルス情報を取得
        elb_tg_health = aws.client('elbv2').describe_target_health(
            TargetGroupArn=elb_tg['TargetGroupArn']
            )['TargetHealthDescriptions']

        # 配下のインスタンス
        for target in elb_tg_health:
            elb_data['Instances'] += [
                instance_infos(
                    target['Target']['Id'],
                    target['TargetHealth'],
                    ec2s)
                ]

        # 収集
        elb_datas += [elb_data]

    return elb_datas

#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#
def collect_elb_datas(elb_list_specified, profile):
    ''' ELBデータ収集 '''
    # 戻り値(list)
    elb_datas = []

    if profile != None:
        from boto3.session import Session
        aws = Session(profile_name=profile)
    else:
        aws = boto3

    # AWS認証情報が正常にセットされているかどうかの確認
    try:
        aws.client('sts').get_caller_identity()
    except:
        # セットされていなかったらヘルプ表示して終了
        print "Cannot access AWS API. Please set AWS credential environments, or use '-p' option.\n"
        os.system("{} -h".format(sys.argv[0]))
        sys.exit()

    # ELBのRAWデータ取得
    elbs_all = aws.client('elb').describe_load_balancers()['LoadBalancerDescriptions']
    elb_name_all = list(map(lambda x: x['LoadBalancerName'], elbs_all))
    elb_name_all = [x['LoadBalancerName'] for x in elbs_all]

    elbv2_lbs_all = aws.client('elbv2').describe_load_balancers()['LoadBalancers']
    elbv2_tgs_all = aws.client('elbv2').describe_target_groups()['TargetGroups']
    elbv2_name_all = list(map(lambda x: x['LoadBalancerName'], elbv2_lbs_all))

    # 対象リストの作成
    if len(elb_list_specified) == 0:
        # 指定が無ければ全てのELBが対象
        elb_list = elb_name_all
        elbv2_list = elbv2_name_all
    else:
        # 引数で指定があれば、指定されたELBのみを対象とする
        elb_list = [x for x in elb_list_specified if x in elb_name_all]
        elbv2_list = [x for x in elb_list_specified if x in elbv2_name_all]

    # EC2の情報取得
    ec2s = aws.client('ec2').describe_instances()['Reservations']

    # ELBデータのサマリ作成
    clb_data(aws, elb_datas, elb_list, elbs_all, ec2s)                          # CLB
    elbv2_data(aws, elb_datas, elbv2_list, elbv2_lbs_all, elbv2_tgs_all, ec2s)  # ALB/NLB

    return elb_datas

#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#
def elb_treeview(elb_datas):
    ''' ツリー図の表示 '''

    #---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#
    def instance_name(name):
        ''' Nameタグがあれば括弧付きで返答 '''
        return ' ({})'.format(name) if name != None else ''

    def status_reason(reason):
        ''' Status_reasonがあったら表示 '''
        return ' - {}'.format(reason) if reason != None else ''

    #---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#

    # タイプごとにループ
    for lb_type in ['classic', 'application', 'network']:
        # タイプに合ったもののみ抽出 + LB名でsort
        elb_datas_per_type = sorted(
            [x for x in elb_datas if x['Type'] == lb_type], key=lambda x: x['LoadBalancerName']
            )

        # 配列が空だったらスキップ
        if len(elb_datas_per_type) == 0:
            continue

        # 最後のELB名を控えておく
        last_lb_name = elb_datas_per_type[-1]['LoadBalancerName']

        print "{}".format(DESCRIPTION[lb_type])
        print " |"

        # ELBごとにループ (ELB名順)
        for elb in elb_datas_per_type:
            # ツリー表記のため縦線の要不要を判定(最後の要素なら縦線不要)
            vertical_line = " " if elb['LoadBalancerName'] == last_lb_name else "|"

            # ELB名・ターゲットグループ名(ALB/NLBのみ)
            print " +- {}".format(elb['LoadBalancerName'])
            if lb_type in ['application', 'network']:
                print " {}  |".format(vertical_line)
                print " {}  {} (TargetGroup)".format(vertical_line, elb['TargetGroupName'])

            # インスタンス
            for instance in sorted(elb['Instances'], key=lambda x: x['Id']):
                print ' {}   +-{}- {} - {:<19}{}{}'.format(
                    vertical_line,
                    MARK.get(instance['Status'], MARK['__default__']),      # ステータスマーク表示
                    instance['AvailabilityZone'],
                    instance['Id'],
                    instance_name(instance['Name']),
                    status_reason(instance['Status_reason'])
                    )

            # 1行あける
            print ' {}'.format(vertical_line)

#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#---#
def main():
    ''' メイン '''

    # 引数処理
    parser = argparse.ArgumentParser(
        description='list ELBs and there target instances in a tree-like (or JSON) format.'
        )
    # ELB名 (.elbnames)
    parser.add_argument('elbnames', metavar='ELB Name', nargs='*',
                        help='specify ELB name(s) you want, or all ELB(s) if not set.')
    # 出力モード (.output)
    parser.add_argument('-o', '--output', metavar='format', default='tree',
                        choices=['tree', 'json', 'raw'],
                        help='output format, json or tree. default is \'tree\'.')
    # プロファイル指定
    parser.add_argument('-p', '--profile', metavar='profile',
                        help='Use a specific profile from your credential file.')

    args = parser.parse_args()

    # ELBデータの収集
    elb_datas = collect_elb_datas(args.elbnames, args.profile)

    # 表示
    if args.output == 'tree':
        # ツリー表示
        elb_treeview(elb_datas)
    elif args.output == 'json':
        # JSON出力
        print json.dumps(elb_datas)
    else:
        # そのまま(raw)
        print elb_datas

if __name__ == '__main__':
    main()