#!/usr/bin/env python3
"""Copy files out of an OCI image archive without mounting the image.

usage: extract-oci-files.py <image.tar> <destination> <path-in-image>...

`buildah mount` needs a `buildah unshare` session, and that fails on the CI host (memfd_create: Invalid argument,
no subuid range for the nested user namespace). Exporting an OCI archive needs no mount, so the files are read from
the layers directly. Later layers win; whiteouts delete earlier files. Each file is written with mode 0755.
"""

from __future__ import annotations

import gzip
import io
import json
import os
import sys
import tarfile
from pathlib import Path


def normalize(name: str) -> str:
    return os.path.normpath(name).lstrip("/")


def blob(archive: tarfile.TarFile, digest: str) -> bytes:
    algorithm, _, value = digest.partition(":")
    member = archive.extractfile(f"blobs/{algorithm}/{value}")
    if member is None:
        raise SystemExit(f"extract-oci-files: blob {digest} is missing from the archive")
    return member.read()


def layers_of(archive: tarfile.TarFile) -> list[dict]:
    index = json.loads(archive.extractfile("index.json").read())
    manifest = None
    for entry in index["manifests"]:
        document = json.loads(blob(archive, entry["digest"]))
        if "layers" in document:
            manifest = document
            break
        # an image index: take its first manifest
        for child in document.get("manifests", []):
            candidate = json.loads(blob(archive, child["digest"]))
            if "layers" in candidate:
                manifest = candidate
                break
        if manifest:
            break
    if manifest is None:
        raise SystemExit("extract-oci-files: no image manifest in the archive")
    return manifest["layers"]


def open_layer(data: bytes, media_type: str) -> tarfile.TarFile:
    if media_type.endswith("zstd"):
        raise SystemExit("extract-oci-files: zstd layers are not supported; export with gzip compression")
    stream = io.BytesIO(gzip.decompress(data) if media_type.endswith("gzip") else data)
    return tarfile.open(fileobj=stream, mode="r:")


def main() -> int:
    if len(sys.argv) < 4:
        print(__doc__, file=sys.stderr)
        return 2
    image, destination = sys.argv[1], Path(sys.argv[2])
    wanted = [normalize(path) for path in sys.argv[3:]]
    found: dict[str, bytes | None] = {}
    with tarfile.open(image) as archive:
        for layer in layers_of(archive):
            with open_layer(blob(archive, layer["digest"]), layer["mediaType"]) as content:
                for member in content:
                    name = normalize(member.name)
                    directory, base = os.path.split(name)
                    if base == ".wh..wh..opq":
                        for path in wanted:
                            if path.startswith(directory + "/"):
                                found[path] = None
                    elif base.startswith(".wh."):
                        removed = os.path.join(directory, base[4:])
                        for path in wanted:
                            if path == removed or path.startswith(removed + "/"):
                                found[path] = None
                    elif name in wanted and member.isfile():
                        found[name] = content.extractfile(member).read()
    missing = [path for path in wanted if found.get(path) is None]
    if missing:
        print("extract-oci-files: not found in the image: " + ", ".join(missing), file=sys.stderr)
        return 1
    destination.mkdir(parents=True, exist_ok=True)
    for path in wanted:
        target = destination / os.path.basename(path)
        target.write_bytes(found[path])
        target.chmod(0o755)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())