#!/usr/bin/env python3
"""oci-push.py: push an image an engine saved to a registry, over the registry's HTTP API.

    oci-push.py <archive.tar> <registry> <repository> <tag> [--platform os/arch]

<archive.tar> is what a container engine's `save` writes for one image (with `--format
oci-archive` where the engine needs it): an OCI image layout (index.json, blobs/sha256/...).
<registry> is the registry's base URL, http:// or https:// (the forge's own address on the private
network, http://<host>:3400, so nothing goes through the public tunnel); <repository> is the
image's name there (`transaction-science/wai-relay`). The manifest for --platform (default
linux/amd64) is chosen from the layout, every blob it names is read from the archive and checked
against its digest, uploaded unless the registry has it, and the manifest is put under <tag>. The
registry is then asked for <tag> and must answer with the manifest's digest. No engine is needed,
so pushing over plain HTTP needs no engine setting.

Credentials: OCI_USER and OCI_TOKEN (a token as the password; a registry that takes tokens reads
no user name from it) answer the registry's challenge, Basic or Bearer. For a Bearer challenge
the token is fetched from the challenge's realm, or from OCI_TOKEN_URL when it is set: the
forge's realm names its public address, and OCI_TOKEN_URL keeps that request on the private
network too. Neither is printed.

Exit 0 when the tag serves the pushed manifest, 1 saying what failed. Standard library only.
"""
import base64
import hashlib
import json
import os
import re
import sys
import tarfile
import urllib.error
import urllib.parse
import urllib.request

INDEX_TYPES = {"application/vnd.oci.image.index.v1+json", "application/vnd.docker.distribution.manifest.list.v2+json"}
MANIFEST_TYPES = {"application/vnd.oci.image.manifest.v1+json", "application/vnd.docker.distribution.manifest.v2+json"}


def die(msg):
    print(f"oci-push: {msg}", file=sys.stderr)
    sys.exit(1)


class Layout:
    """The archive's blobs, by digest."""

    def __init__(self, path):
        self.tar = tarfile.open(path)
        self.members = {m.name.lstrip("./"): m for m in self.tar.getmembers() if m.isfile()}
        if "index.json" not in self.members:
            die(f"{path} holds no index.json: save the image as an OCI image layout")

    def read(self, name):
        return self.tar.extractfile(self.members[name]).read()

    def blob_member(self, digest):
        algo, _, hexd = digest.partition(":")
        name = f"blobs/{algo}/{hexd}"
        if algo != "sha256" or name not in self.members:
            die(f"the archive holds no blob {digest}")
        return self.members[name]

    def blob(self, digest):
        data = self.tar.extractfile(self.blob_member(digest)).read()
        if "sha256:" + hashlib.sha256(data).hexdigest() != digest:
            die(f"the archive's blob {digest} is not that digest")
        return data


def choose(layout, want_os, want_arch):
    """The image manifest for the platform: its descriptor and its bytes."""
    found = []

    def walk(descs, platform):
        for d in descs:
            p = d.get("platform") or platform
            mt = d.get("mediaType")
            if mt in INDEX_TYPES:
                walk(json.loads(layout.blob(d["digest"])).get("manifests", []), p)
            elif mt in MANIFEST_TYPES:
                body = layout.blob(d["digest"])
                if not p:
                    cfg = json.loads(layout.blob(json.loads(body)["config"]["digest"]))
                    p = {"os": cfg.get("os"), "architecture": cfg.get("architecture")}
                if (p.get("os"), p.get("architecture")) == (want_os, want_arch):
                    found.append((d, body))

    walk(json.loads(layout.read("index.json")).get("manifests", []), None)
    digests = {d["digest"] for d, _ in found}
    if len(digests) != 1:
        die(f"the archive holds {len(digests)} manifests for {want_os}/{want_arch}, not one")
    return found[0]


class Registry:
    def __init__(self, base, repo):
        self.base = base.rstrip("/")
        self.repo = repo
        self.auth = None
        self.user = os.environ.get("OCI_USER", "")
        self.secret = os.environ.get("OCI_TOKEN", "")

    def url(self, path):
        return path if re.match(r"https?://", path) else self.base + path

    def _basic(self):
        return "Basic " + base64.b64encode(f"{self.user or 'token'}:{self.secret}".encode()).decode()

    def _answer(self, challenge):
        if not self.secret:
            die("the registry asks for credentials: set OCI_TOKEN")
        scheme, _, params = challenge.partition(" ")
        if scheme.lower() == "basic":
            self.auth = self._basic()
            return
        if scheme.lower() != "bearer":
            die(f"the registry's challenge is {scheme}, neither Basic nor Bearer")
        p = dict(re.findall(r'(\w+)="([^"]*)"', params))
        realm = os.environ.get("OCI_TOKEN_URL") or p.get("realm")
        if not realm:
            die("the registry's Bearer challenge names no realm")
        q = {"service": p.get("service", ""), "scope": f"repository:{self.repo}:pull,push"}
        req = urllib.request.Request(realm + ("&" if "?" in realm else "?") + urllib.parse.urlencode(q),
                                     headers={"Authorization": self._basic()})
        try:
            with urllib.request.urlopen(req, timeout=60) as r:
                t = json.loads(r.read())
        except urllib.error.HTTPError as e:
            die(f"the token endpoint answered {e.code}")
        tok = t.get("token") or t.get("access_token")
        if not tok:
            die("the token endpoint answered no token")
        self.auth = "Bearer " + tok

    def call(self, method, path, body=None, headers=None, ok=(200,), length=None):
        for attempt in (0, 1):
            h = dict(headers or {})
            if self.auth:
                h["Authorization"] = self.auth
            if length is not None:
                h["Content-Length"] = str(length)
            data = body() if callable(body) else body
            req = urllib.request.Request(self.url(path), data=data, method=method, headers=h)
            try:
                with urllib.request.urlopen(req, timeout=600) as r:
                    return r.status, r.headers, r.read()
            except urllib.error.HTTPError as e:
                if e.code == 401 and attempt == 0 and "WWW-Authenticate" in e.headers:
                    self._answer(e.headers["WWW-Authenticate"])
                    continue
                if e.code in ok:
                    return e.code, e.headers, e.read()
                if method == "HEAD" and e.code == 404:
                    return 404, e.headers, b""
                die(f"{method} {path.split('?')[0]}: {e.code} {e.read()[:300]!r}")
            except urllib.error.URLError as e:
                die(f"{method} {self.url(path).split('?')[0]}: {e.reason}")
        die(f"{method} {path}: the registry did not take the credentials")


def main(argv):
    args = [a for a in argv if not a.startswith("--platform")]
    plat = next((a.split("=", 1)[1] for a in argv if a.startswith("--platform=")), None)
    if "--platform" in argv:
        i = argv.index("--platform")
        plat = argv[i + 1]
        args = argv[:i] + argv[i + 2:]
    if len(args) != 4:
        die("usage: oci-push.py <archive.tar> <registry> <repository> <tag> [--platform os/arch]")
    archive, base, repo, tag = args
    want_os, _, want_arch = (plat or "linux/amd64").partition("/")
    layout = Layout(archive)
    desc, manifest = choose(layout, want_os, want_arch)
    m = json.loads(manifest)
    reg = Registry(base, repo)
    blobs = [m["config"]] + m.get("layers", [])
    # Every blob is checked against its digest before anything is sent.
    for b in blobs:
        d = b["digest"]
        h = hashlib.sha256()
        f = layout.tar.extractfile(layout.blob_member(d))
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
        if "sha256:" + h.hexdigest() != d:
            die(f"the archive's blob {d} is not that digest: nothing is pushed")
    for b in blobs:
        d = b["digest"]
        member = layout.blob_member(d)
        if reg.call("HEAD", f"/v2/{repo}/blobs/{d}", ok=(200, 404))[0] == 200:
            print(f"oci-push: {d[:19]} is there already ({member.size} bytes)")
            continue
        status, hdrs, _ = reg.call("POST", f"/v2/{repo}/blobs/uploads/", body=b"", ok=(202,))
        loc = hdrs.get("Location")
        if not loc:
            die("the registry started an upload with no Location")
        loc = urllib.parse.urljoin(reg.base + "/", loc)
        sep = "&" if "?" in loc else "?"
        reg.call("PUT", f"{loc}{sep}digest={urllib.parse.quote(d)}", body=lambda: layout.tar.extractfile(member),
                 headers={"Content-Type": "application/octet-stream"}, ok=(201, 204), length=member.size)
        print(f"oci-push: {d[:19]} uploaded ({member.size} bytes)")
    mt = desc.get("mediaType") or m.get("mediaType")
    digest = "sha256:" + hashlib.sha256(manifest).hexdigest()
    reg.call("PUT", f"/v2/{repo}/manifests/{tag}", body=manifest, headers={"Content-Type": mt}, ok=(201, 204), length=len(manifest))
    _, hdrs, got = reg.call("GET", f"/v2/{repo}/manifests/{tag}", headers={"Accept": mt})
    served = hdrs.get("Docker-Content-Digest") or "sha256:" + hashlib.sha256(got).hexdigest()
    if served != digest:
        die(f"{repo}:{tag} serves {served}, not the pushed {digest}")
    print(f"oci-push: {repo}:{tag} is {digest} ({want_os}/{want_arch}, {len(m.get('layers', []))} layers)")


if __name__ == "__main__":
    main(sys.argv[1:])