#!/usr/bin/env python

# genoci: generate oci images from a yaml spec.  See the README.md for
# example usage.
#
# Copyright (C) 2017 Cisco Inc
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from distutils.version import StrictVersion
import os
import os.path
import shlex
import shutil
import sys
import umoci
import yaml

# by default work in current dir
basedir = os.getcwd()
layoutname="oci"

defconfig = {
        "driver": "vfs",
        "layoutdir": "./oci",
        "btrfsmount": "./btrfs",
        "lvbasedir": "./lvm",
        "vg": "stacker",
        "lvdev": "nbd0"
    }

def parse_config_file(filename):
    configs = {}
    with open(filename, "r") as outfile:
        configs = yaml.load(outfile)
    return configs

def parse_config():
    filename = "./atom_config.yaml"
    retvalues = defconfig
    configs = {}
    if os.path.exists(filename):
        configs = parse_config_file(filename)

    if "driver" in configs:
        retvalues["driver"] = configs["driver"]

    if retvalues["driver"] == "btrfs":
        if "btrfsmount" in configs:
            retvalues["btrfsmount"] = configs["btrfsmount"]
        if "lofile" in configs:
            retvalues["lofile"] = configs["lofile"]
        else:
            retvalues["lofile"] = "btrfs.img"

    elif retvalues["driver"] == "lvm":
        if "lvbasedir" in configs:
            retvalues["lvbasedir"] = configs["lvbasedir"]
        if "vg" in configs:
            retvalues["vg"] = configs["vg"]
        if "lvdev" in configs:
            retvalues["vg"] = configs["vg"]
        if "lofile" in configs:
            retvalues["lofile"] = configs["lofile"]
        else:
            retvalues["lofile"] = "./lvm.img"

    if "layoutdir" in configs:
        retvalues["layoutdir"] = configs["layoutdir"]

    return retvalues

configvalues = parse_config()
basedir, layoutname = os.path.split(configvalues["layoutdir"])

unpackdir = basedir + "/unpacked"

def is_version(x):
    try:
        v = StrictVersion(x)
        return True
    except:
        return False

def do_tag(OCI, tag):
    OCI.DelTag(tag, False)
    OCI.Tag(tag)
    newtag = OCI.NextVersionTag(tag)
    OCI.AddTag(tag, newtag)

# TODO - handle arguments
def run_file(OCI, fname, in_host_ns = False):
    return OCI.RunInChroot(fname, in_host_ns)

def run_shell(OCI, shell, in_host_ns = False):
    return OCI.ShellInChroot(shell, in_host_ns)

OCI = umoci.Umoci(basedir, layoutname, configvalues)

if len(sys.argv) < 2:
    usage()
    sys.exit(1)
if sys.argv[1] == "-h" or sys.argv[1] == "--help":
    usage()
    sys.exit(0)

tags = []
with open(sys.argv[1], "r") as outfile:
    tags = yaml.load(outfile)

todo = []
completed = []
for k in tags:
    todo.append(k)

print("Loaded yaml: %s" % str(tags))
print("To do: %s" % str(todo))

# run: run a command inside the image.
def process_run(OCI, run):
    if os.path.exists(run):
        if run_file(OCI, run) is not True:
            print("Failed running %s for %s" % (run, k))
            return False
    else:
        if run_shell(OCI, run)is not True:
            print("Failed executing shell for %s" % k)
            return False
    return True

def process_runs(OCI, tags, k):
    try:
        runs = tags[k]["run"]
    except:
        return True
    if type(runs) is list:
        for r in runs:
            if not process_run(OCI, r):
                return False
    else:
        if not process_run(OCI, runs):
            return False
    return True

def process_cmds(OCI, tags, k):
    try:
        runs = tags[k]["entrypoint"]
    except:
        return True
    OCI.configs["entrypoint"] = shlex.split(runs)
    return True

def copy_file(OCI, pair):
    copy = pair.split()
    if len(copy) != 2:
        print("Error: badly formed \"from to\" file copy spec: %s" % pair)
        return False
    return OCI.CopyFile(copy[0], copy[1])

# Run pre or post hooks, to do things like create a needed
# mount for a run action.  This runs in the host namespaces.
def run_hook(OCI, hook):
    if os.path.exists(hook):
        if run_file(OCI, hook, in_host_ns = True) is not True:
            print("Failed running %s for %s" % (run, k))
            return False
    else:
        if run_shell(OCI, hook, in_host_ns = True) is not True:
            print("Failed executing shell for %s" % k)
            return False
    return True

def handle_pre_post(OCI, tags, k, which):
    try:
        hooks = tags[k][which]
    except:
        return True
    if type(hooks) is list:
        for h in hooks:
            if not run_hook(OCI, h.replace("%ROOT%", OCI.chrootdir)):
                return False
        return True
    else:
        return run_hook(OCI, hooks.replace("%ROOT%", OCI.chrootdir))

# copy: copy a file from host into image
def copy_files(OCI, tags, k):
    try:
        files = tags[k]["copy"]
    except:
        return True
    if type(files) is list:
        for f in files:
            if not copy_file(OCI, f):
                return False
        return True
    else:
        return copy_file(OCI, files)

def install_pkg(OCI, pkg):
    s = os.path.splitext(pkg)
    if len(s) != 2:
        print("Bad package name: %s" % pkg)
        return False
    pkgfile = OCI.chrootdir + "/" + pkg
    if pkg[0:7] == "http://" or pkg[0:8] == "https://" or pkg[0:6] == "ftp://":
        bname = os.path.basename(pkg)
        if not os.path.exists(bname):
            pkgfile = OCI.chrootdir + "/" + bname
            urllib2.urlopen(pkg, timeout = 5)
            contents= response.read()
            f = open(pkgfile, 'w' )
            f.write(contents)
            f.close()
        pkg = bname

    delfile = False
    if not os.path.exists(OCI.chrootdir + "/" + pkg):
        delfile = True
        shutil.copy(pkg, OCI.chrootdir)

    if s[1] == ".deb":
        cmd = 'chroot %s dpkg -i %s' % (OCI.chrootdir, pkg)
    elif s[1] == ".rpm":
        cmd = 'chroot %s rpm -Uvh %s' % (OCI.chrootdir, pkg)
    else:
        print("Unknown package type: %s" % s[1])
        return False
    print("Installing package: %s" % pkg)
    ret = os.system(cmd)
    if delfile:
        os.remove(pkgfile)
    if ret != 0:
        print("Error installing")
        return False
    print ("%s installed" % pkg)
    return True

# install: install an rpm from the host
# arguments: a list of package filenames, separated by spaces
def install_pkgs(OCI, tags, k):
    try:
        files = tags[k]["install"]
    except:
        return True
    if type(files) is list:
        for pkg in files:
            if not install_pkg(OCI, pkg):
                return False
        return True
    pkglist = files.split()
    for p in pkglist:
        if not install_pkg(OCI, p):
            return False
    return True

def expand_tarball(OCI, tar):
    cmd = 'tar --acls --xattrs --auto-compress -xf %s -C %s' % (tar, OCI.chrootdir)
    print("Expanding tarball: %s" % tar)
    print("Expanding tarball %s into %s" % (tar, OCI.chrootdir))
    ret = os.system(cmd)
    if ret != 0:
        print("Error installing")
        return False
    print("%s expanded" % tar)
    return True

# expand: expand a tarball from the host
def expand_tarballs(OCI, tags, k):
    try:
        files = tags[k]["expand"]
    except:
        return True
    if type(files) is list:
        for p in files:
            if not expand_tarball(OCI, p):
                return False
        return True
    pkglist = files.split()
    for p in pkglist:
        if not expand_tarball(OCI, p):
            return False
    return True

while len(todo) != 0:
    for k in todo:
        # check to see whether it depends on another not-completed base
        OCI.clearconfig()
        base = tags[k]["base"]
        if base in todo:
            continue
        todo.remove(k)

        print("Building %s" % k)
        # expand base
        if base == "empty":
            OCI.Unpack("empty")
        elif OCI.HasTag(base):
            OCI.Unpack(base)
        else:
            print("Unknown base: %s" % base)
            sys.exit(1)


        # Run the conversion step
        if not handle_pre_post(OCI, tags, k, "pre"):
            handle_pre_post(OCI, tags, k, "post")
            print("Failed running pre hooks")
            sys.exit(1)
        if not copy_files(OCI, tags, k):
            print("Failed copying files")
            handle_pre_post(OCI, tags, k, "post")
            sys.exit(1)
        if not install_pkgs(OCI, tags, k):
            print("Failed installing packages")
            handle_pre_post(OCI, tags, k, "post")
            sys.exit(1)
        if not expand_tarballs(OCI, tags, k):
            print("Failed expanding tarballs")
            handle_pre_post(OCI, tags, k, "post")
            sys.exit(1)
        if not process_runs(OCI, tags, k):
            print("Failed running actions")
            handle_pre_post(OCI, tags, k, "post")
            sys.exit(1)
        if not process_cmds(OCI, tags, k):
            print("Failed setting entrypoint")
            handle_pre_post(OCI, tags, k, "post")
            sys.exit(1)
        handle_pre_post(OCI, tags, k, "post")

        do_tag(OCI, k)
        completed.append(k)

if os.path.exists(unpackdir):
    shutil.rmtree(unpackdir)
print("Done")
