import boto3
from pprint import pprint
from utility import have_terraform_tags

CF_LOGGING_BUCKET = "shared-orcd-cloudfront-logs"

def get_cf_info(id):
    client = boto3.client("cloudfront")
    distribution = client.get_distribution(Id=id)["Distribution"]
    tags = client.list_tags_for_resource(Resource=distribution["ARN"])["Tags"]["Items"]
    cf_info = {
        "ARN": distribution["ARN"],
        "HaveTerraformTags": have_terraform_tags(tags),
        "LoggingBucket": distribution["DistributionConfig"]
        .get("Logging", {})
        .get("Bucket", None),
        "LoggingPrefix": distribution["DistributionConfig"]
        .get("Logging", {})
        .get("Prefix", None),
    }
    return cf_info


def list_cf_distributions():
    client = boto3.client("cloudfront")
    next_marker = None
    cf_list = []
    while True:
        if next_marker:
            response = client.list_distributions(Marker=next_marker)
        else:
            response = client.list_distributions()

        for cf in response["DistributionList"]["Items"]:
            cf_list.append(get_cf_info(cf["Id"]))

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

    return cf_list


def verify_cloudfront_logging():
    """Checks cloudfront logging settings to find distributions which
    do not have CF_LOGGING_BUCKET set as logging bucket and/or
    do not have their logging prefix started with 'AWSLogs/account_id'"""

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

    return [
        lb
        for lb in list_cf_distributions()
        if lb["LoggingBucket"] != CF_LOGGING_BUCKET
        or lb["LoggingPrefix"].startswith(f"AWSLogs/{account_id}")
    ]


if __name__ == "__main__":
    pprint(verify_cloudfront_logging())