#!/usr/bin/env python3

import boto3
from botocore.exceptions import ClientError
import argparse


class color:
    BOLD = '\033[1m'
    GREEN = '\033[92m'
    END = '\033[0m'

def getListOfElb(elbv2):
    paginator = elbv2.get_paginator('describe_load_balancers')
    try:
        response_iterator = paginator.paginate(
            PaginationConfig={
                'StartingToken': None
            }
        )
    except ClientError as e:
        if e.response['Error']['Code'] == 'InternalServiceErrorException':
            # An error occurred on the server side.
            # Deal with the exception here, and/or rethrow at your discretion.
            raise e
        elif e.response['Error']['Code'] == 'InvalidParameterException':
            # You provided an invalid value for a parameter.
            # Deal with the exception here, and/or rethrow at your discretion.
            raise e
        elif e.response['Error']['Code'] == 'InvalidRequestException':
            # You provided a parameter value that is not valid for the current state of the resource.
            # Deal with the exception here, and/or rethrow at your discretion.
            raise e
        elif e.response['Error']['Code'] == 'ResourceNotFoundException':
            # We can't find the resource that you asked for.
            # Deal with the exception here, and/or rethrow at your discretion.
            raise e
        elif e.response['Error']['Code'] == 'ExpiredTokenException':
            print('Please re-run your `auth.py`, as token is expired.')
            raise e
        elif e.response['Error']['Code'] == 'ExpiredToken':
            print('Please re-run your `auth.py`, as token is expired.')
            raise e
    else:
        return response_iterator


def getElbAttribute(elbv2, elbName, args, elbType):
    if elbType == 'elb':
        elbAttribute = elbv2.describe_load_balancer_attributes(
            LoadBalancerName=elbName
        )
    else:
        elbAttribute = elbv2.describe_load_balancer_attributes(
            LoadBalancerArn=elbName
        )
    return elbAttribute

def search(atrribute, key):
    valRet = [val for val in atrribute['Attributes'] if val['Key'] == key]
    if valRet:
        return valRet
    else:
        return [{'Value': 'None'}]


def awsConn(args, elbType, region):
    # Define aws session and the connection.
    try:
        session = boto3.Session(profile_name='default',
                                region_name=region)
        elbv2 = session.client(elbType)
    except ClientError as e:
        if e.response['Error']['Code'] == 'ExpiredToken':
            print('Please re-run your `auth.py`, as token is expired.')
    else:
        # Identify type of value looking for
        if elbType == 'elb':
            name = 'LoadBalancerDescriptions'
        else:
            name = 'LoadBalancers'

        # List all the Loabd Balancer for the elb type
        loadBalancers = []
        for page in getListOfElb(elbv2):
            for elb in page[name]:
                loadBalancers.append(elb)

        if args.list:
            print(color.BOLD + color.GREEN + region + ': ' + elbType + color.END)
            for value in loadBalancers:
                if elbType == 'elb':
                    print(
                        f"{value['LoadBalancerName']}: {getElbAttribute(elbv2, value['LoadBalancerName'], args, elbType)['LoadBalancerAttributes']['AccessLog']['Enabled']}")
                else:
                    attribute = getElbAttribute(
                        elbv2, value['LoadBalancerArn'], args, elbType)
                    print('{0}:\n\t{1}: {2}: {3}'.format(value['LoadBalancerName'], search(attribute, 'access_logs.s3.enabled')[0]['Value'], search(
                        attribute, 'access_logs.s3.bucket')[0]['Value'], search(attribute, 'access_logs.s3.prefix')[0]['Value']))
        elif args.enabled:
            print(color.BOLD + color.GREEN + region + ': ' + elbType + color.END)
            for value in loadBalancers:
                if elbType == 'elb':
                    if getElbAttribute(elbv2, value['LoadBalancerName'], args, elbType)['LoadBalancerAttributes']['AccessLog']['Enabled']:
                        print(f"{value['LoadBalancerName']}")
                else:
                    attribute = getElbAttribute(
                        elbv2, value['LoadBalancerArn'], args, elbType)
                    if search(attribute, 'access_logs.s3.enabled')[0]['Value'] == 'true':
                        print('{0}'.format(value['LoadBalancerName']))
        elif args.disabled:
            print(color.BOLD + color.GREEN + region + ': ' + elbType + color.END)
            for value in loadBalancers:
                if elbType == 'elb':
                    if not getElbAttribute(elbv2, value['LoadBalancerName'], args, elbType)['LoadBalancerAttributes']['AccessLog']['Enabled']:
                        print(f"{value['LoadBalancerName']}")
                else:
                    attribute = getElbAttribute(
                        elbv2, value['LoadBalancerArn'], args, elbType)
                    if search(attribute, 'access_logs.s3.enabled')[0]['Value'] == 'false':
                        print('{0}'.format(value['LoadBalancerName']))

def main(args):
    if not (args.list or args.enabled or args.disabled):
        parser.print_help()
    else:
        if not args.region:
            try:
                ec2 = boto3.client('ec2')
            except ClientError as e:
                if e.response['Error']['Code'] == 'ExpiredToken':
                    print('Please re-run your `auth.py`, as token is expired.')
            else:
                regions = [region['RegionName'] for region in ec2.describe_regions()['Regions'] if region['OptInStatus'] == 'opt-in-not-required']
                for elbType in 'elb', 'elbv2':
                    for region in regions:
                        awsConn(args, elbType, region)
        else:
            for elbType in 'elb', 'elbv2':
                awsConn(args, elbType, args.region)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(prog="check_logs", description="List load blanacer which has log enabled / not enabled")
    parser.add_argument(
        '--list', '-l', help="List system only (default)", action="store_true", default=False)
    parser.add_argument('--enabled', '-e', help="List all elb with logs",
                        action="store_true", default=False)
    parser.add_argument('--disabled', '-d', help="List all elb without logs (default)",
                        action="store_true", default=False)
    parser.add_argument(
        '--region', '-r', help="Override search to single regions, default to all regions with no opt-in require.", default=None)
    args = parser.parse_args()

    main(args)
        