#!/usr/bin/env python3
"""Extract one file from a digest-pinned public OCI image without docker (NK-WP-0058-T02). Read-only.

  python3 tools/contract-tests/oci_extract.py --image forgejo.coulomb.social/coulomb/key-cape@sha256:<64 hex> \
      --file keycape --out DIR/keycape

The manifest must hash to the pinned digest, and every layer read must hash to its manifest digest, so the file
is exactly the one in the pinned image. The pull token is the registry's anonymous one; no credential is used.
The newest layer holding the file wins, as in the image's own filesystem. Exit 0 when the file is written.
"""
import argparse
import hashlib
import io
import json
import re
import sys
import tarfile
import urllib.parse
import urllib.request
from pathlib import Path

ACCEPT = ', '.join(('application/vnd.oci.image.manifest.v1+json',
                   'application/vnd.docker.distribution.manifest.v2+json'))


class ExtractError(Exception):
    pass


def get(url, token=None, accept=None, opener=urllib.request.urlopen):
    headers = {**({'Authorization': 'Bearer ' + token} if token else {}), **({'Accept': accept} if accept else {})}
    with opener(urllib.request.Request(url, headers=headers), timeout=120) as response:
        return response.read()


def pull_token(host, repository, opener=urllib.request.urlopen):
    query = urllib.parse.urlencode({'service': 'container_registry', 'scope': f'repository:{repository}:pull'})
    return json.loads(get(f'https://{host}/v2/token?{query}', opener=opener))['token']


def extract(image, wanted, out, opener=urllib.request.urlopen):
    match = re.fullmatch(r'([^/]+)/(.+)@(sha256:[0-9a-f]{64})', image)
    if not match:
        raise ExtractError('image must be host/repository@sha256:<64 hex>')
    host, repository, digest = match.groups()
    token = pull_token(host, repository, opener)
    raw = get(f'https://{host}/v2/{repository}/manifests/{digest}', token, ACCEPT, opener)
    if 'sha256:' + hashlib.sha256(raw).hexdigest() != digest:
        raise ExtractError('manifest_digest_mismatch')
    manifest = json.loads(raw)
    if 'layers' not in manifest:
        raise ExtractError('not_a_single_platform_manifest')
    for layer in reversed(manifest['layers']):
        blob = get(f"https://{host}/v2/{repository}/blobs/{layer['digest']}", token, opener=opener)
        if 'sha256:' + hashlib.sha256(blob).hexdigest() != layer['digest']:
            raise ExtractError('layer_digest_mismatch')
        with tarfile.open(fileobj=io.BytesIO(blob), mode='r:*') as archive:
            for member in archive.getmembers():
                if member.isfile() and member.name.lstrip('./') == wanted.lstrip('/'):
                    out.parent.mkdir(parents=True, exist_ok=True)
                    out.write_bytes(archive.extractfile(member).read())
                    out.chmod(0o755)
                    return {'image': digest, 'layer': layer['digest'], 'file': wanted,
                            'sha256': hashlib.sha256(out.read_bytes()).hexdigest()}
    raise ExtractError('file_not_in_image')


def main(argv=None):
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument('--image', required=True)
    p.add_argument('--file', required=True)
    p.add_argument('--out', required=True, type=Path)
    a = p.parse_args(argv)
    try:
        print(json.dumps(extract(a.image, a.file, a.out)))
    except ExtractError as error:
        print(json.dumps({'status': 'refused', 'reason': str(error)}))
        return 1
    return 0


if __name__ == '__main__':
    sys.exit(main())