import boto3
from pprint import pprint
from utility import get_all_available_regions, have_terraform_tags

LB_LOGGING_BUCKET = "shared-orcd-lb-logs"

def get_load_balancer_info(arn, region):
    elb_client = boto3.client("elbv2", region_name=region)
    attributes_list = elb_client.describe_load_balancer_attributes(LoadBalancerArn=arn)[
        "Attributes"
    ]
    attributes = {}
    for attr in attributes_list:
        attributes[attr["Key"]] = attr["Value"]

    s3_bucket = attributes.get("access_logs.s3.bucket", None)
    s3_prefix = attributes.get("access_logs.s3.prefix", None)

    tags = next(
        (
            res["Tags"]
            for res in elb_client.describe_tags(ResourceArns=[arn])["TagDescriptions"]
            if res["ResourceArn"] == arn
        ),
        [],
    )

    lb_info = {
        "ARN": arn,
        "HaveTerraformTags": have_terraform_tags(tags),
        "LoggingBucket": s3_bucket,
        "LoggingPrefix": s3_prefix,
    }

    return lb_info


def list_load_balancers(region):
    elb_client = boto3.client("elbv2", region_name=region)
    next_marker = None
    lb_list = []
    while True:
        if next_marker:
            response = elb_client.describe_load_balancers(Marker=next_marker)
        else:
            response = elb_client.describe_load_balancers()

        for lb in response["LoadBalancers"]:
            lb_list.append(get_load_balancer_info(lb["LoadBalancerArn"], region))

        next_marker = response.get("NextMarker")
        if not next_marker:
            break

    return lb_list


def verify_lb_logging(regions):
    """Checks load balancer logging settings to find LBs which
    do not have LB_LOGGING_BUCKET set as logging bucket or
    do not have their logging prefix started with account id"""

    complete_lb_list = []
    for region in regions:
        complete_lb_list.extend(list_load_balancers(region))

    sts_client = boto3.client("sts")
    account_id = sts_client.get_caller_identity().get("Account")

    return [
        lb
        for lb in complete_lb_list
        if (
            lb["LoggingBucket"] != f"{LB_LOGGING_BUCKET}-{region}"
            and lb["LoggingBucket"] != LB_LOGGING_BUCKET
        )
        or lb["LoggingPrefix"].startswith(account_id)
    ]

if __name__ == "__main__":
    regions = get_all_available_regions()
    pprint(verify_lb_logging(regions))