#!/usr/bin/python3
"""
Initialize the sysroot and pull and deploy an OStree commit

Initializes a clean ostree based system root, pulls the given `commit` and
creates a deployment from it using `osname` as the new stateroot (see [1]).

Since OStree internally uses a hardlink farm to create the file system tree
for the deployment from the commit data, the mountpoints for the final image
need to be supplied via the `mounts` option, as hardlinks must not span
across file systems and therefore the boundaries need to be known when doing
the deployment.

Creating a deployment also entails generating the Boot Loader Specification
entries to boot the system, which contain this the kernel command line.
The `rootfs` option can be used to indicate the root file system, containing
the sysroot and the deployments. Additional kernel options can be passed via
`kernel_opts`.

[1] https://ostree.readthedocs.io/en/latest/manual/deployment/
"""


import contextlib
import os
import sys
import subprocess

import osbuild.api
import osbuild.sources
from osbuild.util import selinux


SCHEMA = """
"required": ["commit", "osname"],
"properties": {
  "commit": {
    "description": "checksum of the OSTree commit",
    "type": "string"
  },
  "mounts": {
    "description": "Mount points of the final file system",
    "type": "array",
    "items": {
      "description": "Description of one mount point",
      "type": "object",
      "required": ["path"],
      "properties": {
        "path": {
          "description": "The path of the mount point",
          "type": "string"
        },
        "mode": {
          "description": "The mode of the mount point",
          "type": "integer",
          "default": 493
        }
      }
    }
  },
  "osname": {
    "description": "Name of the stateroot to be used in the deployment",
    "type": "string"
  },
  "kernel_opts": {
    "description": "Additional kernel command line options",
    "type": "array",
    "items": {
      "description": "A single kernel command line option",
      "type": "string"
    }
  },
  "ref": {
    "description": "OStree ref to create and use for deployment",
    "type": "string"
  },
  "remotes": {
    "description": "Configure remotes for the system repository",
    "type": "array",
    "items": {
      "description": "Description of a remote",
      "type": "object",
      "required": ["name", "url"],
      "properties": {
        "name": {
          "description": "Identifier for the remote",
          "type": "string"
        },
        "url": {
          "description": "URL of the remote",
          "type": "string"
        },
        "branches": {
          "type": "array",
          "items": {
            "description": "Configured branches for the remote",
            "type": "string"
          }
        },
        "gpgkeys": {
          "type": "array",
          "items": {
            "description": "GPG keys for the remote to verify commits",
            "type": "string"
          }
        }
      }
    }
  },
  "rootfs": {
    "description": "Identifier to locate the root file system",
    "type": "object",
    "oneOf": [{
      "required": ["uuid"]
    }, {
      "required": ["label"]
    }],
    "properties": {
      "label": {
        "description": "Identify the root file system by label",
        "type": "string"
      },
      "uuid": {
        "description": "Identify the root file system by UUID",
        "type": "string"
      }
    }
  },
  "populate_var": {
    "description": "Populate $stateroot/var via systemd-tmpfiles",
    "type": "boolean",
    "default": false
  }
}
"""


def ostree(*args, _input=None, **kwargs):
    args = list(args) + [f'--{k}={v}' for k, v in kwargs.items()]
    print("ostree " + " ".join(args), file=sys.stderr)
    subprocess.run(["ostree"] + args,
                   encoding="utf-8",
                   stdout=sys.stderr,
                   input=_input,
                   check=True)


class MountGuard(contextlib.AbstractContextManager):
    def __init__(self):
        self.mounts = []

    def mount(self, source, target, bind=True, ro=False, mode="0755"):
        options = []
        if bind:
            options += ["bind"]
        if ro:
            options += ["ro"]
        if mode:
            options += [mode]

        args = ["--make-private"]
        if options:
            args += ["-o", ",".join(options)]

        subprocess.run(["mount"] + args + [source, target], check=True)
        self.mounts += [{"source": source, "target": target}]

        subprocess.run(["mount"] + args + [source, target], check=True)

    def unmount(self):

        while self.mounts:
            mount = self.mounts.pop()  # FILO: get the last mount
            target = mount["target"]
            subprocess.run(["umount", "--lazy", target],
                           check=True)

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.unmount()


def make_fs_identifier(desc):
    for key in ["uuid", "label"]:
        val = desc.get(key)
        if val:
            return f"{key.upper()}={val}"
    raise ValueError("unknown rootfs type")


def populate_var(sysroot):
    # Like anaconda[1] and Fedora CoreOS dracut[2]
    # [1] pyanaconda/payload/rpmostreepayload.py
    # [2] ignition-ostree-populate-var.sh

    for target in ('lib', 'log'):
        os.makedirs(f"{sysroot}/var/{target}", exist_ok=True)

    for target in ('home', 'roothome', 'lib/rpm', 'opt', 'srv',
                   'usrlocal', 'mnt', 'media', 'spool', 'spool/mail'):

        if os.path.exists(f"{sysroot}/var/{target}"):
            continue

        res = subprocess.run(["systemd-tmpfiles", "--create", "--boot",
                              "--root=" + sysroot,
                              "--prefix=/var/" + target],
                             encoding="utf-8",
                             stdout=sys.stderr,
                             check=False)

        # According to systemd-tmpfiles(8), the return values are:
        #  0 → success
        # 65 → so some lines had to be ignored, but no other errors
        # 73 → configuration ok, but could not be created
        #  1 → other error
        if res.returncode not in [0, 65]:
            raise RuntimeError(f"Failed to provision /var/{target}")


# pylint: disable=too-many-statements
def main(tree, sources, options):
    commit = options["commit"]
    osname = options["osname"]
    rootfs = options.get("rootfs")
    mounts = options.get("mounts", [])
    kopts = options.get("kernel_opts", [])
    ref = options.get("ref", commit)
    remotes = options.get("remotes", [])
    pop_var = options.get("populate_var", False)

    ostree("admin", "init-fs", "--modern", tree,
           sysroot=tree)

    print(f"Fetching ostree commit {commit}")
    osbuild.sources.get("org.osbuild.ostree", [commit])

    source_repo = f"{sources}/org.osbuild.ostree/repo"

    ostree("pull-local", source_repo, commit,
           repo=f"{tree}/ostree/repo")

    if ref != commit:
        ostree("refs", "--create", ref, commit,
               repo=f"{tree}/ostree/repo")

    ostree("admin", "os-init", osname, sysroot=tree)
    # this created a state root at `osname`
    stateroot = f"{tree}/ostree/deploy/{osname}"

    kargs = []

    if rootfs:
        rootfs_id = make_fs_identifier(rootfs)
        kargs += [f"--karg=root={rootfs_id}"]

    for opt in kopts:
        kargs += [f"--karg-append={opt}"]

    with MountGuard() as mounter:
        for mount in mounts:
            path = mount["path"].lstrip("/")
            path = os.path.join(tree, path)
            os.makedirs(path, exist_ok=True)
            os.chmod(path, mount.get("mode", 0o755))
            mounter.mount(path, path)


        ostree("admin", "deploy", ref,
               *kargs,
               sysroot=tree,
               os=osname)

    # now that we have a deployment, we do have a sysroot
    sysroot = f"{stateroot}/deploy/{commit}.0"

    if pop_var:
        populate_var(stateroot)

    ostree("config", "set", "sysroot.readonly", "true",
           repo=f"{tree}/ostree/repo")

    # deploying a tree creates new files that need to be properly
    # labeled for SELinux. In theory, ostree will take care of
    # this by loading the SELinux config from the deployment and
    # then applying the labels; but it does so conditionally on
    # is_selinux_enabled(2), which in our container is FALSE
    # Therefore we have to do the same dance as ostree does, at
    # least for now, and manually re-label the affected paths
    se_policy = None

    for p in ["etc/selinux", "usr/etc/selinux"]:
        se_home = os.path.join(sysroot, p)
        cfgfile = os.path.join(se_home, "config")
        if not os.path.isfile(cfgfile):
            continue

        with open(cfgfile, 'r') as f:
            cfg = selinux.parse_config(f)
        se_policy = selinux.config_get_policy(cfg)

    if se_policy:
        spec = f"{se_home}/{se_policy}/contexts/files/file_contexts"
        # kernel, initramfs & BLS config snippets were
        # written to {tree}/boot
        selinux.setfiles(spec, tree, "/boot")
        # various config files will be created as a result
        # of the 3-way configuration merge, see ostree(3)
        selinux.setfiles(spec, sysroot, "/etc")
        # if we populated /var, we need to fix its labels
        selinux.setfiles(spec, stateroot, "/var")

    for remote in remotes:
        name = remote["name"]
        url = remote["url"]
        branches = remote.get("branches", [])
        gpgkeys = remote.get("gpgkeys", [])

        extra_args = []
        if not gpgkeys:
            extra_args += ["--no-gpg-verify"]

        ostree("remote", "add",
               "--if-not-exists",
               *extra_args,
               name, url,
               *branches,
               repo=f"{tree}/ostree/repo")

        for key in gpgkeys:
            ostree("remote", "gpg-import", "--stdin", name,
                   repo=f"{tree}/ostree/repo",
                   _input=key)


if __name__ == '__main__':
    stage_args = osbuild.api.arguments()
    r = main(stage_args["tree"],
             stage_args["sources"],
             stage_args["options"])
    sys.exit(r)