lb-ip-edit.py

import boto3
from logging import getLogger, INFO
import os

class LBTargetIpManager:
    """
    ロードバランサーAのターゲットグループに別のロードバランサーBのプライベートIP（ENI）を指定する際に、
    ENIが不定期に変更となる仕様のため、その変更をターゲットに反映させる

    Attributes
    ----------
    aws_access_key_id: str
        defaultのプロファイル以外を使用する場合に指定する。
        Lambdaに本クラスを設定する場合は、IAMロールにて権限を与えるため指定は不要。
    aws_secret_access_key: str
        同上
    region_name: str
        同上
    LB_ARN: str
        ロードバランサーAのARN
    LB_TARGET_GROUP_ARN: sttr
        ロードバランサーAのロードバランサーBを登録したターゲットグループのARN
    TARGET_LB_ARN: str
        ロードバランサーBのARN
    lb_registed_addresses: list
        ロードバランサーAのロードバランサーBを登録したターゲットグループに登録されているIPアドレスのリスト
    target_lb_attached_eni_addresses: list
        ロードバランサーBのENIに割り振られたIPアドレスのリスト
    TARGET_PORT: int
        ロードバランサーAのロードバランサーBに対する通信のポート番号
    logger: RootLogger
        ログ出力
        
    """
    def __init__(self, *, 
                aws_access_key_id='', 
                aws_secret_access_key='', 
                region_name='ap-northeast-1',
                lb_arn,
                lb_target_group_arn,
                target_lb_arn):
        """
        Parameters
        ----------
        aws_access_key_id: str
            defaultのプロファイル以外を使用する場合に指定する。
            Lambdaに本クラスを設定する場合は、IAMロールにて権限を与えるため指定は不要。
        aws_secret_access_key: str
            同上
        region_name: str
            同上
        lb_arn: str
            ロードバランサーAのARN
        lb_target_group_arn: sttr
            ロードバランサーAのロードバランサーBを登録したターゲットグループのARN
        target_lb_arn: str
            ロードバランサーBのARN
        """

        self.session = boto3.Session() if aws_access_key_id == '' or aws_secret_access_key == '' else boto3.session.Session(aws_access_key_id=aws_access_key_id,aws_secret_access_key=aws_secret_access_key,region_name=region_name)
        
        self.LB_ARN = lb_arn
        self.LB_TARGET_GROUP_ARN = lb_target_group_arn
        self.TARGET_LB_ARN       = target_lb_arn

        self.lb_registed_addresses = []
        self.target_lb_attached_eni_addresses = []

        self.TARGET_PORT = 80

        self.logger = getLogger(__name__)
        self.logger.setLevel(INFO)


    def handler(self):
        """
        本クラスの一連の処理を実行するエントリポイント
        """
        try:
            if self.is_lb_active():
                self.logger.info('Load Barancer status is Active !')
                self.describe_lb_registed_addresses()
                self.describe_lb_attached_eni_addresses()
                self.set_target_ip_addresses_from_attached_enis()
            else:
                raise Exception('Error : Load Barancer status is not active or not found !!')
        except Exception as e:
            self.logger.warn(e)
    
    def is_lb_active(self):
        """
        ロードバランサーAが有効状態にあるかを判定する

        Returns
        -------
        is_lb_active : bool
            有効=True, 無効=False
        """
        self.logger.info('<START> is_lb_active()')

        elbv2 = self.session.client('elbv2')
        res = elbv2.describe_load_balancers(LoadBalancerArns=[self.LB_ARN])

        if self.is_api_success(res):
            self.logger.info('<SUCCESS> API elbv2.describe_load_balancers()')

            lbs = res['LoadBalancers']
            return len(lbs) > 0 and lbs[0]['State']['Code'] == 'active'
        else:
            self.api_error(res, 'elbv2.describe_load_balancers() result HTTPStatusCode is not 200') 

    def describe_lb_registed_addresses(self):
        """
        ロードバランサーAのロードバランサーBを登録したターゲットグループに登録されているIPアドレスのリストを照会し
        lb_registed_addressesに格納する
        """
        self.logger.info('<START> describe_lb_registed_addresses()')

        elbv2 = self.session.client('elbv2')

        # LBに指定しているターゲットのIPアドレス情報を取得
        res = elbv2.describe_target_health(
            TargetGroupArn=self.LB_TARGET_GROUP_ARN
        )

        if self.is_api_success(res):
            self.logger.info('<SUCCESS> API elbv2.describe_target_health()')

            # statusがdrainingのものは解除中なので対象外とし、有効なターゲットのみ抽出する
            filter_by_draining = filter(lambda x: x['TargetHealth']['State'] != 'draining', res['TargetHealthDescriptions'])
            self.lb_registed_addresses = list(map(lambda x: x['Target']['Id'], filter_by_draining))
            self.logger.info(f'<Data> lb_registed_addresses : {self.lb_registed_addresses}')
        else:
            self.api_error(res, 'elbv2.describe_target_health() result HTTPStatusCode is not 200')      


    def describe_lb_attached_eni_addresses(self):
        """
        ロードバランサーBのENIに割り振られたIPアドレスのリストを照会し
        target_lb_attached_eni_addressesに格納する
        """
        self.logger.info('<START> describe_lb_attached_eni_addresses()')
        
        ec2 = self.session.client('ec2')

        # LBの「ターゲットに指定しているLB」に付与されたENIの情報を取得
        target_lb_name = self.TARGET_LB_ARN.split('/')[2]
        res = ec2.describe_network_interfaces(
            Filters=[
                {
                    'Name': 'description',
                    'Values': [ f'*{target_lb_name}/*' ]
                },
                {
                    'Name': 'attachment.status',
                    'Values': [ 'attaching', 'attached' ]
                },
            ],
        )
        if self.is_api_success(res):
            self.logger.info('<SUCCESS> API ec2.describe_network_interfaces()')

            # ENIからプライベートIPの情報のみ抽出
            self.target_lb_attached_eni_addresses = list(map(lambda x: x['PrivateIpAddress'], res['NetworkInterfaces']))

            self.logger.info(f'<Data> target_lb_attached_eni_addresses : {self.target_lb_attached_eni_addresses}')
        else:
            self.api_error(res, 'ec2.describe_network_interfaces result HTTPStatusCode is not 200')
 
    def set_target_ip_addresses_from_attached_enis(self):
        """
        ロードバランサーAのロードバランサーBを登録したターゲットグループに登録されているIPアドレスのリストを
        ロードバランサーBのENIのIPアドレスに一致するように登録及び解除を行う
        """
        self.logger.info('<START> set_target_ip_addresses_from_attached_enis()')

        elbv2 = self.session.client('elbv2')

        # 現在のターゲットのIPとENIのIPを比較し、登録及び解除が必要なそれぞれのアドレスのリストを作成
        register_addersses   = list(filter(lambda x: x not in self.lb_registed_addresses, self.target_lb_attached_eni_addresses))
        deregister_addersses = list(filter(lambda x: x not in self.target_lb_attached_eni_addresses, self.lb_registed_addresses))

        self.logger.info(f'<DATA> Register address list is {register_addersses}')
        self.logger.info(f'<DATA> Deregister address list is {deregister_addersses}')

        # 登録の必要がある分を実行
        if len(register_addersses) > 0:
            res = elbv2.register_targets(
                TargetGroupArn=self.LB_TARGET_GROUP_ARN,
                Targets=list(map(lambda address: {'Id': address, 'Port': self.TARGET_PORT} ,register_addersses))
            )

            if self.is_api_success(res):
                self.logger.info('<SUCCESS> API elbv2.register_targets()')
            else:
                self.api_error(res, 'API elbv2.register_targets() result HTTPStatusCode is not 200')
        
        # 解除の必要がある分を実行
        if len(deregister_addersses) > 0:
            res = elbv2.deregister_targets(
                TargetGroupArn=self.LB_TARGET_GROUP_ARN,
                Targets=list(map(lambda address: {'Id': address, 'Port': self.TARGET_PORT} ,deregister_addersses))
            )

            if self.is_api_success(res):
                self.logger.info('<SUCCESS> API elbv2.deregister_targets()')
            else:
                self.api_error(res, 'API elbv2.deregister_targets() result HTTPStatusCode is not 200')
    
    def is_api_success(self, response):
        """
        各APIが正常に成功したかを判定する

        Parameters
        ----------
        response: dict
            APIのレスポンスの辞書データ

        Returns
        -------
        is_api_success: bool
            成功=True, 失敗=False
        """
        return response['ResponseMetadata']['HTTPStatusCode'] == 200
    
    def api_error(self, response, message):
        """
        各APIが失敗した場合に、後続の処理を行いたくないのでエラーを上げる

        Parameters
        ----------
        response: dict
            APIのレスポンスの辞書データ
        message: str
            ログに出力したいエラーメッセージ

        Returns
        -------
        error: Exception
            エラーオブジェクト
        """
        self.logger.warn(response['ResponseMetadata']['HTTPStatusCode'])
        raise Exception(f'Error : {message}')

def lambda_handler(event, context):
    """
    Lambdaがコールするハンドラ

    Parameters
    ----------
    event: dict
        Lambdaから渡されるイベント
    context: LambdaContext
        Lambdaから渡されるコンテキスト

    Returns
    -------
    event: dict
        Lambdaから渡されるイベント
    """
    LBM = LBTargetIpManager(lb_arn=os.environ.get('LB_ARN'),
                            lb_target_group_arn=os.environ.get('LB_TARGET_GROUP_ARN'),
                            target_lb_arn=os.environ.get('TARGET_LB_ARN'))
    LBM.handler()

    return event