#!/usr/bin/python
###############################################################################
#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.    #
#                                                                             #
#  Licensed under the Apache License Version 2.0 (the "License"). You may not #
#  use this file except in compliance with the License. A copy of the License #
#  is located at                                                              #
#                                                                             #
#      http://www.apache.org/licenses/LICENSE-2.0/                                        #
#                                                                             #
#  or in the "license" file accompanying this file. This file is distributed  #
#  on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express #
#  or implied. See the License for the specific language governing permis-    #
#  sions and limitations under the License.                                   #
###############################################################################


from aws_xray_sdk.core import xray_recorder

from ..common.log import get_logger

logger = get_logger(__name__)


@xray_recorder.capture("resolve_artifact_metadata")
def resolve_artifact_metadata(s3_client, s3_bucket_name, prefix):
    # artifact_locaton, artifact_size, sha256
    objects = s3_client.list_objects_v2(Bucket=s3_bucket_name, Prefix=prefix)

    artifact_metadata = []

    for obj in objects.get("Contents", []):
        object_key = obj["Key"]
        object_size = obj["Size"]
        file_name_with_extension = object_key.split("/")[-1]

        filename_parts = object_key.split("_")

        is_sha_256 = False

        for part in filename_parts:
            is_sha_256 = part.find("sha256") != -1

        if not is_sha_256:
            # get filename without extension and add _sha256.txt
            filename_parts = file_name_with_extension.split(".")
            filename = (
                ".".join(filename_parts[0:-1])
                if len(filename_parts) > 1
                else ".".join(filename_parts)
            )
            sha_value = None

            if [
                obj["Key"]
                for obj in objects.get("Contents", [])
                if obj["Key"] == f"{prefix}/{filename}_sha256.txt"
            ]:
                sha_response = s3_client.get_object(
                    Bucket=s3_bucket_name,
                    Key=f"{prefix}/{filename}_sha256.txt",
                )
                sha_value = (
                    sha_response["Body"].read().decode("utf-8")
                    if sha_response
                    else "SHA value was not generated by SSM investigation document"
                )

            artifact_metadata.append(
                {
                    "artifact_location": object_key,
                    "artifact_size": object_size,
                    "sha256": sha_value,
                }
            )

    return artifact_metadata