#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
#  Authors:
#        Sam Thursfield <sam.thursfield@codethink.co.uk>
#        Chandan Singh <csingh43@bloomberg.net>

#
# This plugin was originally developped in the https://gitlab.com/BuildStream/bst-plugins-container/
# repository and was copied from 192aa75a62161b655b051d5b69f17d9233cf4dfe
#

"""
docker - stage files from Docker images
=======================================

**Usage:**

.. code:: yaml

   # Specify the docker source kind
   kind: docker

   # Specify the registry endpoint, defaults to Docker Hub (optional)
   registry-url: https://registry.hub.docker.com

   # Image path
   image: library/alpine

   # Alternatively, you can specify a full url, the registry endpoint
   # is assumed to be at the root of the domain
   # url: https://registry.hub.docker.com/library/alpine

   # Image tag to follow (optional)
   track: latest

   # Specify the digest of the exact image to use (required)
   ref: 6c9f6f68a131ec6381da82f2bff978083ed7f4f7991d931bfa767b7965ebc94b

   # Some images are built for multiple platforms. When tracking a tag, we
   # will choose which image to use based on these settings. Default values
   # are chosen based on the output of `uname -m` and `uname -s`, but you
   # can override them.
   #
   # architecture: arm64
   # os: linux
   #
   # **Since**: 2.0.1

   # Specify the version to be reported as the *guess_version* when reporting
   # SourceInfo
   #
   # Since: 2.5
   #
   version: 1.2

Note that Docker images may contain device nodes. BuildStream elements cannot
contain device nodes so those will be dropped. Any regular files in the /dev
directory will also be dropped.

See `built-in functionality doumentation
<https://docs.buildstream.build/master/buildstream.source.html#core-source-builtins>`_ for
details on common configuration options for sources.


Reporting `SourceInfo <https://docs.buildstream.build/master/buildstream.source.html#buildstream.source.SourceInfo>`_
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The docker source reports the URL of the docker registry as the *url*.

Further, the docker source reports the `SourceInfoMedium.OCI_IMAGE
<https://docs.buildstream.build/master/buildstream.source.html#buildstream.source.SourceInfoMedium.OCI_IMAGE>`_
*medium* and the `SourceVersionType.OCI_DIGEST
<https://docs.buildstream.build/master/buildstream.source.html#buildstream.source.SourceVersionType.OCI_DIGEST>`_
*version_type*, for which it reports the content digest of the docker image as the *version*.

Additionally, the docker source reports the docker image name through
the ``image-name`` key of the *extra_data*.

As such, after removing the scheme from the URL (i.e. remove ``https://``) the same docker
image can be obtained by calling:

.. code:: bash

   docker pull <url-without-scheme>/<image-name>@<version>

Since the docker source does not have any way to guess what tag is associated
to the digest, or what release version that would mean; the docker source exposes
the ``version`` configuration attribute to allow explicit specification of the
*guess_version*.
"""

import hashlib
import json
import os
import platform
import shutil
import tarfile
import urllib.parse

from collections import OrderedDict

import requests

from buildstream import Source, SourceError
from buildstream.utils import (
    save_file_atomic,
    sha256sum,
    link_files,
    move_atomic,
    get_umask,
)

#
# Soft import of buildstream symbols only available in newer versions
#
# The BST_MIN_VERSION will provide a better user experience.
#
try:
    from buildstream import SourceInfoMedium, SourceVersionType
except ImportError:
    pass

_DOCKER_HUB_URL = "https://registry.hub.docker.com"


def parse_bearer_authorization_challenge(text):
    # Hand-written and probably broken parsing of the Www-Authenticate
    # response. I can't find a built-in way to parse this, but I probably
    # didn't look hard enough.
    if not text.startswith("Bearer "):
        raise SourceError("Unexpected Www-Authenticate response: %{}".format(text))

    pairs = {}
    text = text[len("Bearer ") :]
    for pair in text.split(","):
        key, value = pair.split("=")
        pairs[key] = value[1:-1]
    return pairs


def default_architecture():
    machine = platform.machine()
    if machine == "x86_64":
        return "amd64"
    elif machine == "aarch64":
        return "arm64"
    else:
        return machine


def default_os():
    return platform.system().lower()


# Variant of urllib.parse.urljoin() allowing multiple path components at once.
def urljoin(url, *args):
    for arg in args:
        if not url.endswith("/"):
            url += "/"
        url = urllib.parse.urljoin(url, arg.lstrip("/"))
    return url


# Handles authentication with a bearer token
class BearerAuth(requests.auth.AuthBase):
    def __init__(self, api_timeout=3):
        self.token = None
        self.api_timeout = api_timeout

    def __call__(self, r):
        if self.token:
            r.headers["Authorization"] = "Bearer {}".format(self.token)
        return r

    def refresh_token(self, auth_challenge):
        auth_vars = parse_bearer_authorization_challenge(auth_challenge)
        # Respond to an Www-Authenticate challenge by requesting the necessary
        # token from the 'realm' (endpoint) that we were given in the challenge.
        request_url = "{realm}?service={service}&scope={scope}".format(**auth_vars)
        response = requests.get(request_url, timeout=self.api_timeout)
        response.raise_for_status()
        self.token = response.json()["token"]


# DockerManifestError
#
# Raised if something goes wrong while querying an image manifest from a remote
# registry.
#
class DockerManifestError(SourceError):
    def __init__(self, message, manifest=None):
        super().__init__(message)
        self.manifest = manifest


class DockerRegistryV2Client:
    def __init__(self, endpoint, api_timeout=3):
        self.endpoint = endpoint
        self.api_timeout = api_timeout

        self.auth = BearerAuth(api_timeout)

    def _request(self, subpath, extra_headers=None, stream=False, _reauthorized=False):
        if not extra_headers:
            extra_headers = {}

        headers = {"content-type": "application/json"}
        headers.update(extra_headers)

        url = urljoin(self.endpoint, "v2", subpath)
        response = requests.get(url, headers=headers, stream=stream, timeout=self.api_timeout, auth=self.auth)

        if response.status_code == requests.codes["unauthorized"] and not _reauthorized:
            # This request requires (re)authorization. See:
            # https://docs.docker.com/registry/spec/auth/token/
            auth_challenge = response.headers["Www-Authenticate"]
            self.auth.refresh_token(auth_challenge)
            return self._request(subpath, extra_headers=extra_headers, _reauthorized=True)
        else:
            response.raise_for_status()

            return response

    # digest():
    #
    # Calculate a Docker-compatible digest of an arbitrary string of bytes.
    #
    # Args:
    #    content (bytes): Content to hash
    #
    # Returns:
    #    (str) A Docker-compatible digest of 'content'
    @staticmethod
    def digest(content):
        digest_hash = hashlib.sha256()
        digest_hash.update(content)
        return "sha256:" + digest_hash.hexdigest()

    # manifest():
    #
    # Fetches the image manifest for a given image from the remote registry.
    #
    # If this is a "fat" (multiplatform) image, the 'artitecture' and 'os'
    # parameters control which of the available images is chosen.
    #
    # The manifest is returned verbatim, so you need to parse it yourself
    # with json.loads() to get at its contents. The verbatim text can be used
    # to recalculate the content digest, just encode it and pass to .digest().
    # If we returned only the parsed JSON data you wouldn't be able to do this.
    #
    # Args:
    #    image_path (str): Relative path to the image, e.g. library/alpine
    #    reference (str): Either a tag name (such as 'latest') or the content
    #                     digest of an exact version of the image.
    #    architecture (str): Architecture name (amd64, arm64, etc.)
    #    os_ (str): OS name (e.g. linux)
    #
    # Raises:
    #    requests.RequestException, if network errors occur
    #
    # Returns:
    #    (str, str): A tuple of the manifest content as text, and its content hash
    def manifest(
        self,
        image_path,
        reference,
        architecture,
        os_,
    ):
        # pylint: disable=too-many-locals

        accept_types = [
            "application/vnd.docker.distribution.manifest.v2+json",
            "application/vnd.docker.distribution.manifest.list.v2+json",
        ]

        manifest_url = urljoin(image_path, "manifests", urllib.parse.quote(reference))
        response = self._request(manifest_url, extra_headers={"Accept": ",".join(accept_types)})

        try:
            manifest = json.loads(response.text)
        except json.JSONDecodeError as e:
            raise DockerManifestError(
                "Server did not return a valid manifest: {}".format(e),
                manifest=response.text,
            ) from e

        schema_version = manifest.get("schemaVersion")
        if schema_version == 1:
            raise DockerManifestError("Schema version 1 is unsupported.", manifest=response.text)
        if schema_version is None:
            raise DockerManifestError(
                "Manifest did not include the schemaVersion key.",
                manifest=response.text,
            )

        our_digest = self.digest(response.text.encode("utf8"))
        their_digest = response.headers.get("Docker-Content-Digest")

        if not their_digest:
            raise DockerManifestError(
                "Server did not set the Docker-Content-Digest header.",
                manifest=response.text,
            )
        if our_digest != their_digest:
            raise DockerManifestError(
                "Server returned a non-matching content digest. "
                "Our digest: {}, their digest: {}".format(our_digest, their_digest),
                manifest=response.text,
            )

        if manifest["mediaType"] == "application/vnd.docker.distribution.manifest.list.v2+json":
            # This is a "fat manifest", we need to narrow down to a specific
            # architecture.
            for sub in manifest["manifests"]:
                if sub["platform"]["architecture"] == architecture and sub["platform"]["os"]:
                    sub_digest = sub["digest"]
                    return self.manifest(
                        image_path,
                        sub_digest,
                        architecture=architecture,
                        os_=os_,
                    )
            raise DockerManifestError(
                "No images found for architecture {}, OS {}".format(architecture, os_),
                manifest=response.text,
            )
        if manifest["mediaType"] == "application/vnd.docker.distribution.manifest.v2+json":
            return response.text, our_digest
        else:
            raise DockerManifestError(
                "Unsupported manifest type {}".format(manifest["mediaType"]),
                manifest=response.text,
            )

    # blob():
    #
    # Fetch a blob from the remote registry. This is used for getting each
    # layer of an image in tar.gz format.
    #
    # Raises:
    #    requests.RequestException, if network errors occur
    #
    # Args:
    #    image_path (str): Relative path to the image, e.g. library/alpine
    #    blob_digest (str): Content hash of the blob.
    #    download_to (str): Path to a file where the content will be written.
    def blob(self, image_path, blob_digest, download_to):
        blob_url = urljoin(image_path, "blobs", urllib.parse.quote(blob_digest))

        response = self._request(blob_url, stream=True)

        with save_file_atomic(download_to, "wb") as f:
            shutil.copyfileobj(response.raw, f)


class ReadableTarInfo(tarfile.TarInfo):
    """
    The goal is to override `TarFile`'s `extractall` semantics by ensuring that on extraction, the
    files are readable by the owner of the file. This is done by overriding the accessor for the
    `mode` attribute in `TarInfo`, the class that encapsulates the internal meta-data of the tarball,
    so that the owner-read bit is always set.
    """

    # https://github.com/python/mypy/issues/4125
    @property  # type: ignore
    def mode(self):
        # Respect umask instead of the file mode stored in the archive.
        # The only bit used from the embedded mode is the executable bit for files.
        umask = get_umask()
        if self.isdir() or bool(self.__permission & 0o100):
            return 0o777 & ~umask
        else:
            return 0o666 & ~umask

    @mode.setter
    def mode(self, permission):
        self.__permission = permission  # pylint: disable=attribute-defined-outside-init


class DockerSource(Source):
    # pylint: disable=too-many-instance-attributes

    BST_MIN_VERSION = "2.5"

    # Docker identifies images by a content digest calculated from the image's
    # manifest. This corresponds well with the concept of a 'ref' in
    # BuildStream. However, Docker theoretically supports multiple hash
    # methods while BuildStream does not. Right now every Docker registry
    # uses sha256 so let's ignore that issue for the time being.
    @staticmethod
    def _digest_to_ref(digest):
        if digest.startswith("sha256:"):
            return digest[len("sha256:") :]
        else:
            method = digest.split(":")[0]
            raise SourceError("Unsupported digest method: {}".format(method))

    @staticmethod
    def _ref_to_digest(ref):
        return "sha256:" + ref

    def configure(self, node):
        # url is deprecated, but accept it as a valid key so that we can raise
        # a nicer warning.
        node.validate_keys(
            Source.COMMON_CONFIG_KEYS
            + ["architecture", "registry-url", "image", "os", "ref", "track", "url", "version"]
        )

        if "url" in node and ("image" in node or "registry-url" in node):
            raise SourceError("{}: May specify either 'url', or 'image' and 'registry-url'".format(self))

        if "url" not in node and "image" not in node:
            raise SourceError("{}: Must define either 'url' or 'image'".format(self))

        if "url" in node:
            self.url = node.get_str("url")
            translated_url = self.translate_url(self.url)
            scheme, netloc, path, _, _ = urllib.parse.urlsplit(translated_url)
            self.registry_url = urllib.parse.urlunsplit((scheme, netloc, "", "", ""))
            self.image = path

        if "image" in node:
            self.url = None
            self.image = node.get_str("image")
            self.original_registry_url = node.get_str("registry-url", _DOCKER_HUB_URL)
            self.registry_url = self.translate_url(self.original_registry_url)

        self.tag = node.get_str("track", "") or None

        self.architecture = node.get_str("architecture", "") or default_architecture()
        self.os = node.get_str("os", "") or default_os()
        self.version = node.get_str("version", None)

        self.digest = None
        self.load_ref(node)

        self.client = DockerRegistryV2Client(self.registry_url)

        self.manifest = None

    def load_ref(self, node):
        ref = node.get_str("ref", None)
        if ref is not None:
            self.digest = self._ref_to_digest(ref)

    def preflight(self):
        return

    def get_unique_key(self):
        if self.url is not None:
            unique_key = [self.url, self.digest]
        else:
            unique_key = [self.original_registry_url, self.image, self.digest]

        # Backwards compatible method of supporting configuration
        # attributes which affect SourceInfo generation.
        if self.version is not None:
            unique_key.append(self.version)

        return unique_key

    def get_ref(self):
        return None if self.digest is None else self._digest_to_ref(self.digest)

    def set_ref(self, ref, node):
        node["ref"] = ref
        self.digest = self._ref_to_digest(ref)

    def track(self):
        # pylint: disable=arguments-differ

        # If the tracking ref is not specified it's not an error, just silently return
        if not self.tag:
            return None

        with self.timed_activity(
            "Fetching image manifest for image: '{}:{}' from: {}".format(self.image, self.tag, self.registry_url)
        ):
            try:
                _, digest = self.client.manifest(self.image, self.tag, self.architecture, self.os)
            except DockerManifestError as e:
                self.log("Problem downloading manifest", detail=e.manifest)
                raise
            except (OSError, requests.RequestException) as e:
                raise SourceError(e) from e

        return self._digest_to_ref(digest)

    def is_resolved(self):
        return self.digest is not None

    def is_cached(self):
        mirror_dir = self.get_mirror_directory()
        try:
            manifest = self._load_manifest()

            for layer in manifest["layers"]:
                layer_digest = layer["digest"]
                blob_path = os.path.join(mirror_dir, layer_digest + ".tar.gz")
                try:
                    self._verify_blob(blob_path, expected_digest=layer_digest)
                except FileNotFoundError:
                    # digest fetched, but some layer blob not fetched
                    return False
            return True
        except (FileNotFoundError, SourceError):
            return False

    def _load_manifest(self):
        manifest_file = os.path.join(self.get_mirror_directory(), self.digest + ".manifest.json")

        with open(manifest_file, "rb") as f:
            text = f.read()

        real_digest = self.client.digest(text)
        if real_digest != self.digest:
            raise SourceError("Manifest {} is corrupt; got content hash of {}".format(manifest_file, real_digest))

        return json.loads(text.decode("utf-8"))

    def _save_manifest(self, text, path):
        manifest_file = os.path.join(path, self.digest + ".manifest.json")
        with save_file_atomic(manifest_file, "wb") as f:
            f.write(text.encode("utf-8"))

    @staticmethod
    def _verify_blob(path, expected_digest):
        blob_digest = "sha256:" + sha256sum(path)
        if expected_digest != blob_digest:
            raise SourceError("Blob {} is corrupt; got content hash of {}.".format(path, blob_digest))

    def fetch(self):
        # pylint: disable=arguments-differ

        with self.timed_activity(
            "Fetching image {}:{} with digest {}".format(self.image, self.tag, self.digest),
            silent_nested=True,
        ):
            with self.tempdir() as tmpdir:
                # move all files to a tmpdir
                try:
                    manifest = self._load_manifest()
                except FileNotFoundError as e:
                    try:
                        manifest_text, digest = self.client.manifest(
                            self.image,
                            self.digest,
                            self.architecture,
                            self.os,
                        )
                    except requests.RequestException as ee:
                        raise SourceError(ee) from ee

                    if digest != self.digest:
                        raise SourceError(
                            "Requested image {}, got manifest with digest {}".format(self.digest, digest)
                        ) from e
                    self._save_manifest(manifest_text, tmpdir)
                    manifest = json.loads(manifest_text)
                except DockerManifestError as e:
                    self.log("Unexpected manifest", detail=e.manifest)
                    raise
                except (OSError, requests.RequestException) as e:
                    raise SourceError(e) from e

                for layer in manifest["layers"]:
                    if layer["mediaType"] != "application/vnd.docker.image.rootfs.diff.tar.gzip":
                        raise SourceError("Unsupported layer type: {}".format(layer["mediaType"]))

                    layer_digest = layer["digest"]
                    blob_path = os.path.join(tmpdir, layer_digest + ".tar.gz")

                    if not os.path.exists(blob_path):
                        try:
                            self.client.blob(self.image, layer_digest, download_to=blob_path)
                        except (OSError, requests.RequestException) as e:
                            if os.path.exists(blob_path):
                                shutil.rmtree(blob_path)
                            raise SourceError(e) from e

                    self._verify_blob(blob_path, expected_digest=layer_digest)

                # Only if all sources are successfully fetched, move files to staging directory

                # As both the manifest and blobs are content addressable, we can optimize space by having
                # a flat mirror directory. We check one-by-one if there is any need to copy a file out of the tmpdir.
                for fetched_file in os.listdir(tmpdir):
                    move_atomic(
                        os.path.join(tmpdir, fetched_file),
                        os.path.join(self.get_mirror_directory(), fetched_file),
                    )

    def stage(self, directory):
        mirror_dir = self.get_mirror_directory()

        try:
            manifest = self._load_manifest()
        except (OSError, SourceError) as e:
            raise SourceError("Unable to load manifest: {}".format(e)) from e

        try:
            layer_members = []

            # Create a list of members to extract from each layer
            for layer in manifest["layers"]:
                layer_digest = layer["digest"]
                blob_path = os.path.join(mirror_dir, layer_digest + ".tar.gz")

                self._verify_blob(blob_path, expected_digest=layer_digest)

                members_dict, whiteout_paths, opaque_whiteout_dirs = self._get_members_and_whiteout_paths(blob_path)

                # Process opaque whiteouts: remove members under the corresponding directories from previous layers
                for _, prev_members_dict in layer_members:
                    members_to_remove = []
                    for member_path in prev_members_dict:
                        # Check if this member should be removed by any opaque whiteout
                        for opaque_whiteout_dir in opaque_whiteout_dirs:
                            if member_path.startswith(opaque_whiteout_dir):
                                members_to_remove.append(member_path)
                                break
                    for member_path in members_to_remove:
                        del prev_members_dict[member_path]

                # Process whiteouts: remove corresponding members from previous layers
                for _, prev_members_dict in layer_members:
                    for whiteout_path in whiteout_paths:
                        prev_members_dict.pop(whiteout_path, None)

                layer_members.append((blob_path, members_dict))

            # Extract files for each layer
            for blob_path, members_dict in layer_members:
                if not members_dict:
                    # No files to extract from this layer
                    continue

                # Extract files for the current layer
                with tarfile.open(blob_path, tarinfo=ReadableTarInfo) as tar:
                    with self.tempdir() as td:
                        members = list(members_dict.values())

                        if hasattr(tarfile, "tar_filter"):
                            # Python 3.12+ (and older versions with backports)
                            tar.extraction_filter = tarfile.tar_filter
                        else:
                            for member in members:
                                self._assert_tarinfo_safe(member, td)
                        tar.extractall(path=td, members=members)
                        link_files(td, directory)

        except (OSError, SourceError, tarfile.TarError) as e:
            raise SourceError("{}: Error staging source: {}".format(self, e)) from e

    def collect_source_info(self):

        # If the image was configured with "url" only rather
        # than "registry-url" and "image", then self.image
        # at this point will have a leading forward slash "/".
        #
        # Lets just normalize that problem here, as changing
        # the self.image itself can break cache keys.
        #
        image_name = self.image.lstrip("/")

        return [
            self.create_source_info(
                self.registry_url,
                SourceInfoMedium.OCI_IMAGE,
                SourceVersionType.OCI_DIGEST,
                self.digest,
                version_guess=self.version,
                extra_data={"image-name": image_name},
            )
        ]

    @staticmethod
    def _get_members_and_whiteout_paths(layer_tar_path):
        """Return the set of files to remove and extract for a given layer

        :param layer_tar_path: The path to the layer tar.gz file
        :return: Tuple of (members_dict, whiteout_paths, opaque_whiteout_dirs)
          - members_dict: OrderedDict of TarInfo members to extract into staging directory
          - whiteout_paths: list of file paths that should be removed from previous layers
          - opaque_whiteout_dirs: list of directory paths whose contents should be removed from previous layers

        """

        def strip_wh(white_out_file):
            """Strip the prefixing .wh. for given file

            :param white_out_file: path of file
            :return: path without white-out prefix
            """
            # whiteout files have the syntax of `*/.wh.*`
            file_name = os.path.basename(white_out_file)
            path = os.path.join(os.path.dirname(white_out_file), file_name.split(".wh.", maxsplit=1)[1])
            return path

        def is_regular_file(info):
            """Check if file is a non-device file

            :param info: tar member metadata
            :return: if the file is a non-device file
            """
            return not (info.name.startswith("dev/") or info.isdev())

        with tarfile.open(layer_tar_path, tarinfo=ReadableTarInfo) as tar:
            members_dict = OrderedDict()
            whiteout_paths = []
            opaque_whiteout_dirs = []

            for member in tar.getmembers():
                member_basename = os.path.basename(member.name)
                if member_basename == ".wh..wh..opq":
                    # Opaque whiteout file
                    opaque_whiteout_dir = os.path.dirname(member.name)
                    if opaque_whiteout_dir != "":
                        # Add "/" to facilitate path prefix matching
                        opaque_whiteout_dir += "/"
                    opaque_whiteout_dirs.append(opaque_whiteout_dir)
                elif member_basename.startswith(".wh."):
                    # Whiteout file
                    whiteout_paths.append(strip_wh(member.name))
                elif is_regular_file(member):
                    members_dict[member.name] = member

        return members_dict, whiteout_paths, opaque_whiteout_dirs

    # Assert that a tarfile is safe to extract; specifically, make
    # sure that we don't do anything outside of the target
    # directory (this is possible, if, say, someone engineered a
    # tarfile to contain paths that start with ..).
    def _assert_tarinfo_safe(self, member: tarfile.TarInfo, target_dir: str):
        final_path = os.path.abspath(os.path.join(target_dir, member.path))
        if not final_path.startswith(target_dir):
            raise SourceError(
                "{}: Tarfile attempts to extract outside the staging area: "
                "{} -> {}".format(self, member.path, final_path)
            )

        if member.islnk():
            linked_path = os.path.abspath(os.path.join(target_dir, member.linkname))
            if not linked_path.startswith(target_dir):
                raise SourceError(
                    "{}: Tarfile attempts to hardlink outside the staging area: "
                    "{} -> {}".format(self, member.path, final_path)
                )

        # Don't need to worry about symlinks because they're just
        # files here and won't be able to do much harm once we are
        # in a sandbox.


# Plugin entry point
def setup():
    return DockerSource
