#!/bin/sh
# Import a digest-pinned GCP owner image into the session k3d node as linux/amd64.
# Docker 29 `image save --platform linux/amd64` cannot export some OCI indexes
# (and `save` without --platform writes a stub tar without layers). Resolve the
# linux/amd64 child digest from the pinned index, pull that digest into k3d
# with ctr --platform linux/amd64, then prove the local tag is that amd64
# manifest — not merely that a save command exited 0.
set -eu

usage() {
  echo "usage: $0 <pinned-image@sha256:...> <k3d-cluster> <local-image>" >&2
  exit 2
}

[ "$#" -eq 3 ] || usage
if [ -z "${DOCKER_CONFIG:-}" ]; then
  echo "DOCKER_CONFIG is required so manifest inspect uses the login that pulled the pin" >&2
  exit 2
fi

image=$1
cluster=$2
local_image=$3
node="k3d-${cluster}-server-0"
lock="/tmp/cp9-${cluster}-image-import.lock"
auth=
# Per-invocation path so EXIT cleanup cannot delete another import's copy.
remote_pass=/tmp/cp9-gcr.pass.$$
cleanup() {
  if [ -n "$auth" ]; then
    rm -f "$auth"
  fi
  if [ -n "$remote_pass" ]; then
    docker exec "$node" rm -f "$remote_pass" >/dev/null 2>&1 || true
  fi
}
trap cleanup EXIT
# dash does not run EXIT on SIGINT/SIGTERM; exit from these traps so EXIT runs.
trap 'exit 130' INT
trap 'exit 143' TERM

case "$image" in
  *@sha256:*) ;;
  *)
    echo "GCP owner image must be digest pinned: $image" >&2
    exit 2
    ;;
esac

registry=${image%%/*}
case "$registry" in
  gcr.io|*.gcr.io|*.pkg.dev)
    python3 - "$registry" <<'PY'
import re, sys, urllib.parse

registry = sys.argv[1]
parsed = urllib.parse.urlsplit("https://" + registry)
labels = registry.lower().split(".")
label = re.compile(r"[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?\Z")
if (
    parsed.netloc != registry
    or parsed.hostname != registry.lower()
    or parsed.port is not None
    or any(not label.fullmatch(part) for part in labels)
):
    raise SystemExit("GCP registry host is invalid: " + registry)
PY
    token=$(curl --connect-timeout 5 --max-time 20 --fail --silent --show-error \
      --header 'Metadata-Flavor: Google' \
      http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token \
      | python3 -c 'import json,sys; print(json.load(sys.stdin)["access_token"])')
    [ -n "$token" ]
    auth=$(mktemp)
    printf '%s' "$token" >"$auth"
    chmod 600 "$auth"
    unset token
    ;;
esac

# Docker 29 inspect can reject a pin whose bytes still match. Fetch the pin
# independently, require sha256(body)==pin, then use a single manifest or the
# linux/amd64 child after hashing that child too. Do not treat inspect error
# strings as a pin fallback.
resolved=$(python3 - "$image" "$auth" <<'PY'
import hashlib, json, re, subprocess, sys
import urllib.error, urllib.parse, urllib.request

image, auth_path = sys.argv[1], sys.argv[2]
name, pin = image.rsplit("@", 1)
if not pin.startswith("sha256:"):
    sys.exit("GCP owner image must be digest pinned: " + image)
last = name.split("/")[-1]
if ":" in last:
    name = name.rsplit(":", 1)[0]
token = open(auth_path, encoding="utf-8").read() if auth_path else ""
host, repo = name.split("/", 1)
try:
    parsed_host = urllib.parse.urlsplit("https://" + host)
    invalid_host = (
        parsed_host.netloc != host
        or parsed_host.hostname is None
        or parsed_host.username is not None
        or parsed_host.password is not None
    )
    parsed_host.port
except ValueError:
    invalid_host = True
if invalid_host:
    sys.exit("image registry host is invalid: " + host)

class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None

opener = urllib.request.build_opener(NoRedirect())

def pick_amd64(manifest):
    if "manifests" not in manifest:
        return pin
    found = None
    for entry in manifest["manifests"]:
        platform = entry.get("platform") or {}
        media = entry.get("mediaType") or ""
        if platform.get("os") != "linux" or platform.get("architecture") != "amd64":
            continue
        if platform.get("variant"):
            continue
        if "manifest" not in media or "in-toto" in media:
            continue
        found = entry.get("digest")
        break
    if not found:
        sys.exit("index has no linux/amd64 image manifest: " + image)
    return found

def registry_get(digest, accept):
    headers = {"Accept": accept}
    if token:
        headers["Authorization"] = "Bearer " + token
    url = "https://%s/v2/%s/manifests/%s" % (host, repo, digest)
    try:
        with opener.open(urllib.request.Request(url, headers=headers), timeout=30) as resp:
            return resp.read()
    except urllib.error.HTTPError as err:
        if err.code == 401 and not token:
            challenge = challenge_token(err.headers.get("WWW-Authenticate"))
            if challenge:
                headers["Authorization"] = "Bearer " + challenge
                try:
                    with opener.open(urllib.request.Request(url, headers=headers), timeout=30) as resp:
                        return resp.read()
                except urllib.error.HTTPError as retry_err:
                    err = retry_err
        if err.code in (401, 403):
            sys.exit("registry denied the pinned image")
        if 300 <= err.code < 400:
            sys.exit("registry redirected the manifest request")
        sys.exit("registry fetch of the pinned image failed")
    except Exception:
        sys.exit("registry fetch of the pinned image failed")

def require_digest(raw, want, what):
    got = "sha256:" + hashlib.sha256(raw).hexdigest()
    if got != want:
        sys.exit("%s digest does not match fetched manifest" % what)
    return json.loads(raw)

def emit(digest):
    print(name + "@" + digest)
    print(digest)

def challenge_fields(value):
    match = re.fullmatch(r"\s*Bearer\s+(.+)\s*", value or "", re.IGNORECASE)
    if not match:
        return None
    fields = {}
    pattern = re.compile(r'''([A-Za-z][A-Za-z0-9_-]*)\s*=\s*(?:"((?:\\.|[^"])*)"|([^,\s]+))''')
    for item in pattern.finditer(match.group(1)):
        value = item.group(2) if item.group(2) is not None else item.group(3)
        fields[item.group(1).lower()] = value.replace('\\"', '"').replace('\\\\', '\\')
    return fields if fields.get("realm") else None

def challenge_token(value):
    fields = challenge_fields(value)
    if not fields:
        return None
    parsed = urllib.parse.urlsplit(fields["realm"])
    if parsed.scheme != "https" or parsed.netloc != host:
        sys.exit("registry authentication challenge is not same-origin HTTPS")
    query = dict(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True))
    for key in ("service", "scope", "account"):
        if key in fields:
            query[key] = fields[key]
    token_url = urllib.parse.urlunsplit(
        (parsed.scheme, parsed.netloc, parsed.path, urllib.parse.urlencode(query), parsed.fragment)
    )
    request = urllib.request.Request(token_url, headers={"Accept": "application/json"})
    try:
        with opener.open(request, timeout=30) as response:
            payload = json.loads(response.read())
    except urllib.error.HTTPError as err:
        if 300 <= err.code < 400:
            sys.exit("registry authentication challenge redirected")
        sys.exit("registry authentication challenge failed")
    except Exception:
        sys.exit("registry authentication challenge failed")
    bearer = payload.get("token") or payload.get("access_token")
    if not isinstance(bearer, str) or not bearer:
        sys.exit("registry authentication challenge returned no token")
    return bearer

accept_index = (
    "application/vnd.oci.image.index.v1+json,"
    "application/vnd.docker.distribution.manifest.list.v2+json,"
    "application/vnd.oci.image.manifest.v1+json,"
    "application/vnd.docker.distribution.manifest.v2+json"
)
accept_manifest = (
    "application/vnd.oci.image.manifest.v1+json,"
    "application/vnd.docker.distribution.manifest.v2+json"
)

inspect = subprocess.run(
    ["docker", "manifest", "inspect", image],
    capture_output=True,
    text=True,
)
manifest = None
if inspect.returncode == 0:
    manifest = json.loads(inspect.stdout)
else:
    blob = (inspect.stdout or "") + (inspect.stderr or "")
    if "denied" in blob.lower() or "unauthenticated" in blob.lower():
        sys.exit("docker manifest inspect failed for the pinned image")
    manifest = require_digest(registry_get(pin, accept_index), pin, "pinned")

child = pick_amd64(manifest)
if child != pin:
    require_digest(registry_get(child, accept_manifest), child, "linux/amd64")
emit(child)
PY
)
amd64_ref=$(printf '%s\n' "$resolved" | sed -n '1p')
amd64_digest=$(printf '%s\n' "$resolved" | sed -n '2p')

# A Google token never appears in subprocess argv. Public registries pull
# anonymously; authenticated ctr reads its copied credential inside the node.
flock "$lock" python3 - "$node" "$amd64_ref" "$local_image" "$amd64_digest" "$auth" "$remote_pass" <<'PY'
import subprocess, sys

node, amd64_ref, local_image, amd64_digest, auth, remote_pass = sys.argv[1:]

def listed_line():
    listed = subprocess.check_output(
        ["docker", "exec", node, "ctr", "-n", "k8s.io", "images", "ls"],
        text=True,
    )
    matched = [line for line in listed.splitlines() if local_image in line]
    if len(matched) != 1:
        return None
    line = matched[0]
    if amd64_digest in line and "linux/amd64" in line and "linux/arm64" not in line:
        return line
    return None

already = listed_line()
if already is None:
    try:
        if auth:
            subprocess.run(["docker", "cp", auth, node + ":" + remote_pass], check=True)
            pull = (
                "chmod 600 \"$1\"; pass=$(cat \"$1\"); rm -f \"$1\"; "
                "exec ctr -n k8s.io images pull --platform linux/amd64 "
                "--user \"oauth2accesstoken:${pass}\" \"$2\""
            )
            command = ["docker", "exec", node, "sh", "-c", pull, "sh", remote_pass, amd64_ref]
        else:
            command = ["docker", "exec", node, "ctr", "-n", "k8s.io", "images", "pull", "--platform", "linux/amd64", amd64_ref]
        subprocess.run(command, check=True)
    except subprocess.CalledProcessError:
        sys.exit("ctr pull of the linux/amd64 digest failed")
    subprocess.run(
        ["docker", "exec", node, "ctr", "-n", "k8s.io", "images", "tag", amd64_ref, local_image],
        check=True,
    )
    already = listed_line()
if already is None:
    sys.exit("imported %s is not linux/amd64 digest %s" % (local_image, amd64_digest))
print(already)
PY