#!/usr/bin/python3

"""

Assemble an OCI image archive

Assemble an Open Container Initiative[1] image[2] archive, i.e. a

tarball whose contents is in the OCI image layout.

Currently the only required options are `filename` and `architecture`.

The execution parameters for the image, which then should form the base

for the container, can be given via `config`. They have the same format

as the `config` option for the "OCI Image Configuration" (see [2]),

except those that map to the "Go type map[string]struct{}", which are

represented as array of strings.

The final resulting tarball, aka a "orci-archive", can be imported via

podman[3] with `podman pull oci-archive:<archive>`.

[1] https://www.opencontainers.org/

[2] https://github.com/opencontainers/image-spec/

[3] https://podman.io/

"""

import datetime

import json

import os

import subprocess

import sys

import tempfile

DEFAULT_PATH = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"

SCHEMA = """

"additionalProperties": false,

"required": ["architecture", "filename"],

"properties": {

  "architecture": {

    "description": "The CPU architecture of the image",

    "type": "string"

  },

  "filename": {

    "description": "Resulting image filename",

    "type": "string"

  },

  "config": {

    "description": "The execution parameters",

    "type": "object",

    "additionalProperties": false,

    "properties": {

      "Cmd": {

        "type": "array",

        "default": ["sh"],

        "items": {

          "type": "string"

        }

      },

      "Env": {

        "type": "array",

        "default": ["PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"],

        "items": {

          "type": "string"

        }

      },

      "ExposedPorts": {

        "type": "array",

        "items": {

          "type": "string"

        }

      },

      "User": {

        "type": "string"

      },

      "Labels": {

        "type": "object",

        "additionalProperties": true

      },

      "StopSiganl": {

        "type": "string"

      },

      "Volumes": {

        "type": "array",

        "items": {

          "type": "string"

        }

      },

      "WorkingDir": {

        "type": "string"

      }

    }

  }

}

"""

MEDIA_TYPES = {

    "layer": "application/vnd.oci.image.layer.v1.tar",

    "manifest": "application/vnd.oci.image.manifest.v1+json",

    "config": "application/vnd.oci.image.config.v1+json"

}

def sha256sum(path: str) -> str:

    ret = subprocess.run(["sha256sum", path],

                         stdout=subprocess.PIPE,

                         encoding="utf-8",

                         check=True)

    return ret.stdout.strip().split(" ")[0]

def blobs_add_file(blobs: str, path: str, mtype: str):

    digest = sha256sum(path)

    size = os.stat(path).st_size

    os.rename(path, os.path.join(blobs, digest))

    info = {

        "digest": "sha256:" + digest,

        "size": size,

        "mediaType": MEDIA_TYPES[mtype]

    }

    print(f"blobs: +{mtype} ({size}, {digest})")

    return info

def blobs_add_json(blobs: str, js: str, mtype: str):

    js_file = os.path.join(blobs, "temporary.js")

    with open(js_file, "w") as f:

        json.dump(js, f)

    return blobs_add_file(blobs, js_file, mtype)

def blobs_add_layer(blobs: str, tree: str):

    compression = "gzip"

    layer_file = os.path.join(blobs, "layer.tar")

    command = [

        "tar",

        "--selinux",

        "--acls",

        "--xattrs",

        "-cf", layer_file,

        "-C", tree,

    ] + os.listdir(tree)

    print("creating layer")

    subprocess.run(command,

                   stdout=subprocess.DEVNULL,

                   check=True)

    digest = "sha256:" + sha256sum(layer_file)

    print("compressing layer")

    suffix = ".compressed"

    subprocess.run([compression,

                    "-S", suffix,

                    layer_file],

                   stdout=subprocess.DEVNULL,

                   check=True)

    layer_file += suffix

    info = blobs_add_file(blobs, layer_file, "layer")

    info["mediaType"] += compression

    return digest, info

def config_from_options(options):

    command = options.get("Cmd", ["sh"])

    env = options.get("Env", ["PATH=" + DEFAULT_PATH])

    config = {

        "Env": env,

        "Cmd": command

    }

    for name in ["User", "Labels", "StopSignal", "WorkingDir"]:

        item = options.get(name)

        if item:

            config[name] = item

    for name in ["ExposedPorts", "Volumes"]:

        item = options.get(name)

        if item:

            config[name] = {x: {} for x in item}

    print(config)

    return config

def create_oci_dir(tree, output_dir, options):

    architecture = options["architecture"]

    config = {

        "created": datetime.datetime.utcnow().isoformat() + "Z",

        "architecture": architecture,

        "os": "linux",

        "config": config_from_options(options["config"]),

        "rootfs": {

            "type": "layers",

            "diff_ids": []

        }

    }

    manifest = {

        "schemaVersion": 2,

        "config": None,

        "layers": []

    }

    index = {

        "schemaVersion": 2,

        "manifests": []

    }

    blobs = os.path.join(output_dir, "blobs", "sha256")

    os.makedirs(blobs)

    ## layers / rootfs

    digest, info = blobs_add_layer(blobs, tree)

    config["rootfs"]["diff_ids"] = [digest]

    manifest["layers"].append(info)

    ## write config

    info = blobs_add_json(blobs, config, "config")

    manifest["config"] = info

    # manifest

    info = blobs_add_json(blobs, manifest, "manifest")

    index["manifests"].append(info)

    # index

    print("writing index")

    with open(os.path.join(output_dir, "index.json"), "w") as f:

        json.dump(index, f)

    # oci-layout tag

    with open(os.path.join(output_dir, "oci-layout"), "w") as f:

        json.dump({"imageLayoutVersion": "1.0.0"}, f)

def main(tree, output_dir, options):

    filename = options["filename"]

    with tempfile.TemporaryDirectory(dir=output_dir) as tmpdir:

        workdir = os.path.join(tmpdir, "output")

        os.makedirs(workdir)

        create_oci_dir(tree, workdir, options)

        command = [

            "tar",

            "--remove-files",

            "-cf", os.path.join(output_dir, filename),

            f"--directory={workdir}",

        ] + os.listdir(workdir)

        print("creating final archive")

        subprocess.run(command,

                       stdout=subprocess.DEVNULL,

                       check=True)

if __name__ == '__main__':

    args = json.load(sys.stdin)

    r = main(args["tree"], args["output_dir"], args["options"])

    sys.exit(r)