#! /usr/bin/env bash

# Copyright (c) 2018-2023, NVIDIA CORPORATION. All rights reserved.
#
# 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.

if [ ${BASH_VERSION:0:1} -lt 4 ] || [ ${BASH_VERSION:0:1} -eq 4 -a ${BASH_VERSION:2:1} -lt 2 ]; then
    printf "Unsupported %s version: %s\n" "${BASH}" "${BASH_VERSION}" >&2
    exit 1
fi

set -euo pipefail
shopt -s lastpipe

# ---- BEGIN CONFIG ----

config::export() {
    [ -n "${3-}" ] && export "$1=$2" || export "$1=${!1:-$2}"
}

config::load() {
    local -r file="$1" overwrite="${2-}"

    if [ -s "${file}" ] && ! "${BASH}" -n "${file}" > /dev/null 2>&1; then
        printf "Invalid configuration file: %s\n" "${file}"
        exit 1
    fi
    if [ -e "${file}" ]; then
        while IFS=$' \t=' read -r key value; do
            if [[ "${key}" =~ ^ENROOT_[[:upper:]_]+$ ]] ||
               [[ "${key}" =~ ^(MELLANOX_|NVIDIA_) ]] ||
               [[ "${key}" =~ ^SSL_CERT_(DIR|FILE)$ ]] ||
               [[ "${key}" =~ ^(all|no|http|https)_proxy$ ]]; then
                config::export "${key}" "$(eval echo "${value}")" "${overwrite}"
            fi
        done < "${file}"
   fi
}

config::init() {
    local -r file="$1" overwrite="${2-}"

    config::load "${file}" "${overwrite}"
    if [ -d "${file}.d" ]; then
      shopt -s nullglob
      for dropin in "${file}.d"/*".conf"; do
        config::load "${dropin}" "${overwrite}"
      done
      shopt -u nullglob
    fi
}

config::fini() {
    for var in $(compgen -e "ENROOT_"); do
        if [[ "${!var}" =~ ^(no?|N[oO]?|[fF](alse)?|FALSE)$ ]]; then
            unset "${var}"
        fi
    done
}

config::init "/etc/enroot/enroot.conf"

config::export XDG_DATA_HOME           "${HOME:-$(echo ~)}/.local/share"
config::export XDG_CONFIG_HOME         "${HOME:-$(echo ~)}/.config"
config::export XDG_CACHE_HOME          "${HOME:-$(echo ~)}/.cache"
config::export XDG_RUNTIME_DIR         "/run"

config::export ENROOT_LIBRARY_PATH     "/usr/lib/enroot"
config::export ENROOT_SYSCONF_PATH     "/etc/enroot"
config::export ENROOT_RUNTIME_PATH     "${XDG_RUNTIME_DIR}/enroot"
config::export ENROOT_CONFIG_PATH      "${XDG_CONFIG_HOME}/enroot"
config::export ENROOT_CACHE_PATH       "${XDG_CACHE_HOME}/enroot"
config::export ENROOT_DATA_PATH        "${XDG_DATA_HOME}/enroot"
config::export ENROOT_TEMP_PATH        "${TMPDIR:-/tmp}"

config::export ENROOT_GZIP_PROGRAM     "$(command -v pigz > /dev/null && echo pigz || echo gzip)"
config::export ENROOT_ZSTD_OPTIONS     "-1"
config::export ENROOT_SQUASH_OPTIONS   "-comp lzo -noD -exit-on-error"
config::export ENROOT_MAX_PROCESSORS   "$(nproc)"
config::export ENROOT_MAX_CONNECTIONS  10
config::export ENROOT_CONNECT_TIMEOUT  30
config::export ENROOT_TRANSFER_TIMEOUT 0
config::export ENROOT_TRANSFER_RETRIES 0
config::export ENROOT_LOGIN_SHELL      true
config::export ENROOT_ALLOW_SUPERUSER  false
config::export ENROOT_ALLOW_HTTP       false
config::export ENROOT_ROOTFS_WRITABLE  false
config::export ENROOT_NATIVE_OVERLAYFS true
config::export ENROOT_REMAP_ROOT       false
config::export ENROOT_BUNDLE_ALL       false
config::export ENROOT_BUNDLE_CHECKSUM  false
config::export ENROOT_FORCE_OVERRIDE   false

config::fini

(umask 077 && mkdir -p "${ENROOT_CACHE_PATH}" "${ENROOT_DATA_PATH}" "${ENROOT_RUNTIME_PATH}")

# ----- END CONFIG -----

export ENROOT_VERSION="4.1.1"

source "${ENROOT_LIBRARY_PATH}/common.sh"
source "${ENROOT_LIBRARY_PATH}/docker.sh"
source "${ENROOT_LIBRARY_PATH}/runtime.sh"

enroot::usage() {
    case "$1" in
    batch)
        cat <<- EOF
		Usage: ${0##*/} batch [options] [--] CONFIG [COMMAND] [ARG...]
		
		Shorthand version of "${0##*/} start -c CONFIG" where the root filesystem is
		taken from the configuration file using the special directive ENROOT_ROOTFS.
		EOF
        ;;
    bundle)
        cat <<- EOF
		Usage: ${0##*/} bundle [options] [--] IMAGE
		
		Create a self-extracting bundle from a container image.
		
		 Options:
		   -a, --all            Include runtime and user configuration files in the bundle
		   -c, --checksum       Generate an embedded checksum
		   -d, --desc TEXT      Provide a description of the bundle
		   -o, --output BUNDLE  Name of the output bundle file (defaults to "IMAGE.run")
		   -t, --target DIR     Target directory used by --keep (defaults to "\$PWD/BUNDLE")
		   -f, --force          Overwrite an existing bundle
		EOF
        ;;
    create)
        cat <<- EOF
		Usage: ${0##*/} create [options] [--] IMAGE
		
		Create a container root filesystem from a container image.
		
		 Options:
		   -n, --name   Name of the container (defaults to "IMAGE")
		   -f, --force  Overwrite an existing root filesystem
		EOF
        ;;
    exec)
        cat <<- EOF
		Usage: ${0##*/} exec [options] [--] PID COMMAND [ARG...]
		
		Execute a command inside an existing container.
		
		 Options:
		   -e, --env KEY[=VAL]  Export an environment variable inside the container
		EOF
        ;;
    export)
        cat <<- EOF
		Usage: ${0##*/} export [options] [--] NAME
		
		Create a container image from a container root filesystem.
		
		 Options:
		   -o, --output  Name of the output image file (defaults to "NAME.sqsh")
		   -f, --force   Overwrite an existing container image
		EOF
        ;;
    import)
        cat <<- EOF
		Usage: ${0##*/} import [options] [--] URI
		
		Import a container image from a specific location.
		
		 Schemes:
		   docker://[USER@][REGISTRY#]IMAGE[:TAG]  Import a Docker image from a registry
		   dockerd://IMAGE[:TAG]                   Import a Docker image from the Docker daemon
		   podman://IMAGE[:TAG]                    Import a Docker image from a local podman repository

		 Options:
		   -a, --arch    Architecture of the image (defaults to host architecture)
		   -o, --output  Name of the output image file (defaults to "URI.sqsh")
		EOF
        ;;
    digest)
        cat <<- EOF
		Usage: ${0##*/} digest [options] [--] URI
		
		Print the SHA256 digest of a container image.
		
		 Schemes:
		   docker://[USER@][REGISTRY#]IMAGE[:TAG]  Get digest of a Docker image from a registry

		 Options:
		   -a, --arch    Architecture of the image (defaults to host architecture)
		EOF
        ;;
    load)
        cat <<- EOF
		Usage: ${0##*/} load [options] [--] URI
		
		Load a container root filesystem directly from a container image.
		
		 Schemes:
		   docker://[USER@][REGISTRY#]IMAGE[:TAG]  Load a Docker image from a registry

		 Options:
		   -a, --arch    Architecture of the image (defaults to host architecture)
		   -n, --name    Name of the container (defaults to "URI")
		   -f, --force   Overwrite an existing root filesystem
		EOF
        ;;
    list)
        cat <<- EOF
		Usage: ${0##*/} list [options]
		
		List all the container root filesystems on the system.
		
		 Options:
		   -f, --fancy  Display more information in tabular format
		EOF
        ;;
    remove)
        cat <<- EOF
		Usage: ${0##*/} remove [options] [--] NAME...
		
		Delete one or multiple container root filesystems.
		
		 Options:
		   -f, --force  Do not prompt for confirmation
		EOF
        ;;
    start)
        cat <<- EOF
		Usage: ${0##*/} start [options] [--] NAME|IMAGE [COMMAND] [ARG...]
		
		Start a container and invoke the command script within its root filesystem.
		Command and arguments are passed to the script as input parameters.
		
		In the absence of a command script and if a command was given, it will be executed directly.
		Otherwise, an interactive shell will be started within the container.
		
		 Options:
		   -c, --conf CONFIG    Specify a configuration script to run before the container starts
		   -e, --env KEY[=VAL]  Export an environment variable inside the container
		       --rc SCRIPT      Override the command script inside the container
		   -r, --root           Ask to be remapped to root inside the container
		   -w, --rw             Make the container root filesystem writable
		   -m, --mount FSTAB    Perform a mount from the host inside the container (colon-separated)
		EOF
        ;;
    info)
        cat <<- EOF
		Usage: ${0##*/} info
		
		Display enroot configuration information.
		EOF
        ;;
    help|*)
        cat <<- EOF
		Usage: ${0##*/} COMMAND [ARG...]
		
		Command line utility for manipulating container sandboxes.
		
		 Commands:
		   batch  [options] [--] CONFIG [COMMAND] [ARG...]
		   bundle [options] [--] IMAGE
		   create [options] [--] IMAGE
		   digest [options] [--] URI
		   exec   [options] [--] PID COMMAND [ARG...]
		   export [options] [--] NAME
		   import [options] [--] URI
		   list   [options]
		   load   [options] [--] URI
		   remove [options] [--] NAME...
		   start  [options] [--] NAME|IMAGE [COMMAND] [ARG...]
		   info
		   version
		EOF
        ;;
    esac
    exit "$2"
}

enroot::version() {
    printf "%s\n" "${ENROOT_VERSION}"
}

enroot::info() {
    while [ $# -gt 0 ]; do
        case "$1" in
        -h|--help)
            enroot::usage info 0 ;;
        --)
            shift; break ;;
        -?*)
            enroot::usage info 1 ;;
        *)
            break ;;
        esac
    done
    if [ $# -ne 0 ]; then
        enroot::usage info 1
    fi

    printf "Enroot version: %s\n" "${ENROOT_VERSION}"
    printf "Config file: %s\n" "/etc/enroot/enroot.conf"
    printf "\nEnvironment variables:\n"

    for var in $(compgen -e "ENROOT_" | sort); do
        printf "  %s=%s\n" "${var}" "${!var}"
    done
}

enroot::digest() {
    local uri= arch=

    while [ $# -gt 0 ]; do
        case "$1" in
        -a|--arch)
            [ -z "${2-}" ] && enroot::usage digest 1
            arch="$2"
            shift 2
            ;;
        --arch=*)
           [ -z "${1#*=}" ] && enroot::usage digest 1
           arch="${1#*=}"
           shift
           ;;
        -h|--help)
            enroot::usage digest 0 ;;
        --)
            shift; break ;;
        -?*)
            enroot::usage digest 1 ;;
        *)
            break ;;
        esac
    done
    if [ $# -ne 1 ]; then
        enroot::usage digest 1
    fi
    uri="$1"

    runtime::digest "${uri}" "${arch}"
}

enroot::import() {
    local uri= filename= arch=

    while [ $# -gt 0 ]; do
        case "$1" in
        -a|--arch)
            [ -z "${2-}" ] && enroot::usage import 1
            arch="$2"
            shift 2
            ;;
        --arch=*)
           [ -z "${1#*=}" ] && enroot::usage import 1
           arch="${1#*=}"
           shift
           ;;
        -o|--output)
            [ -z "${2-}" ] && enroot::usage import 1
            filename="$2"
            shift 2
            ;;
        --output=*)
           [ -z "${1#*=}" ] && enroot::usage import 1
           filename="${1#*=}"
           shift
           ;;
        -h|--help)
            enroot::usage import 0 ;;
        --)
            shift; break ;;
        -?*)
            enroot::usage import 1 ;;
        *)
            break ;;
        esac
    done
    if [ $# -ne 1 ]; then
        enroot::usage import 1
    fi
    uri="$1"

    runtime::import "${uri}" "${filename}" "${arch}"
}

enroot::load() {
    local uri= name= arch=

    while [ $# -gt 0 ]; do
        case "$1" in
        -a|--arch)
            [ -z "${2-}" ] && enroot::usage load 1
            arch="$2"
            shift 2
            ;;
        --arch=*)
           [ -z "${1#*=}" ] && enroot::usage load 1
           arch="${1#*=}"
           shift
           ;;
        -n|--name)
            [ -z "${2-}" ] && enroot::usage load 1
            name="$2"
            shift 2
            ;;
        --name=*)
            [ -z "${1#*=}" ] && enroot::usage load 1
            name="${1#*=}"
            shift
            ;;
        -f|--force)
            export ENROOT_FORCE_OVERRIDE=y
            shift
            ;;
        -h|--help)
            enroot::usage load 0 ;;
        --)
            shift; break ;;
        -?*)
            enroot::usage load 1 ;;
        *)
            break ;;
        esac
    done
    if [ $# -ne 1 ]; then
        enroot::usage load 1
    fi
    uri="$1"

    runtime::load "${uri}" "${name}" "${arch}"
}

enroot::export() {
    local name= filename=

    while [ $# -gt 0 ]; do
        case "$1" in
        -f|--force)
            export ENROOT_FORCE_OVERRIDE=y
            shift
            ;;
        -o|--output)
            [ -z "${2-}" ] && enroot::usage export 1
            filename="$2"
            shift 2
            ;;
        --output=*)
            [ -z "${1#*=}" ] && enroot::usage export 1
            filename="${1#*=}"
            shift
            ;;
        -h|--help)
            enroot::usage export 0 ;;
        --)
            shift; break ;;
        -?*)
            enroot::usage export 1 ;;
        *)
            break ;;
        esac
    done
    if [ $# -ne 1 ]; then
        enroot::usage export 1
    fi
    name="$1"

    runtime::export "${name}" "${filename}"
}

enroot::create() {
    local image= name=

    while [ $# -gt 0 ]; do
        case "$1" in
        -f|--force)
            export ENROOT_FORCE_OVERRIDE=y
            shift
            ;;
        -n|--name)
            [ -z "${2-}" ] && enroot::usage create 1
            name="$2"
            shift 2
            ;;
        --name=*)
            [ -z "${1#*=}" ] && enroot::usage create 1
            name="${1#*=}"
            shift
            ;;
        -h|--help)
            enroot::usage create 0 ;;
        --)
            shift; break ;;
        -?*)
            enroot::usage create 1 ;;
        *)
            break ;;
        esac
    done
    if [ $# -ne 1 ]; then
        enroot::usage create 1
    fi
    image="$1"

    runtime::create "${image}" "${name}"
}

enroot::start() {
    local name= conf= rc= mounts=() environ=()

    while [ $# -gt 0 ]; do
        case "$1" in
        -c|--conf)
            [ -z "${2-}" ] && enroot::usage start 1
            conf="$2"
            shift 2
            ;;
        --conf=*)
            [ -z "${1#*=}" ] && enroot::usage start 1
            conf="${1#*=}"
            shift
            ;;
        -m|--mount)
            [ -z "${2-}" ] && enroot::usage start 1
            mounts+=("$2")
            shift 2
            ;;
        --mount=*)
            [ -z "${1#*=}" ] && enroot::usage start 1
            mounts+=("${1#*=}")
            shift
            ;;
        -e|--env)
            [ -z "${2-}" ] && enroot::usage start 1
            environ+=("$2")
            shift 2
            ;;
        --env=*)
            [ -z "${1#*=}" ] && enroot::usage start 1
            environ+=("${1#*=}")
            shift
            ;;
        -r|--root)
            root=y
            shift
            ;;
        --rc)
            [ -z "${2-}" ] && enroot::usage start 1
            rc="$2"
            shift 2
            ;;
        --rc=*)
            [ -z "${1#*=}" ] && enroot::usage start 1
            rc="${1#*=}"
            shift
            ;;
        -w|--rw)
            rw=y
            shift
            ;;
        -h|--help)
            enroot::usage start 0 ;;
        --)
            shift; break ;;
        -?*)
            enroot::usage start 1 ;;
        *)
            break ;;
        esac
    done
    if [ $# -lt 1 ]; then
        enroot::usage start 1
    fi
    name="$1"
    shift

    # Check for #ENROOT directives in the container configuration.
    if [ -n "${conf}" ] && [ -f "${conf}" ]; then
        common::checkcmd sed
        config::init <(sed -n '/^#[[:space:]]*ENROOT_/s/#//p' "${conf}") true; config::fini
    fi
    if [ -v root ]; then
        export ENROOT_REMAP_ROOT=y
    fi
    if [ -v rw ]; then
        export ENROOT_ROOTFS_WRITABLE=y
    fi

    runtime::start "${name}" "${rc}" "${conf}"     \
      "$(IFS=$'\n'; echo ${mounts[*]+"${mounts[*]}"})"  \
      "$(IFS=$'\n'; echo ${environ[*]+"${environ[*]}"})" \
      "$@"
}

enroot::exec() {
    local pid= environ=()

    while [ $# -gt 0 ]; do
        case "$1" in
        -e|--env)
            [ -z "${2-}" ] && enroot::usage exec 1
            environ+=("$2")
            shift 2
            ;;
        --env=*)
            [ -z "${1#*=}" ] && enroot::usage exec 1
            environ+=("${1#*=}")
            shift
            ;;
        -h|--help)
            enroot::usage exec 0 ;;
        --)
            shift; break ;;
        -?*)
            enroot::usage exec 1 ;;
        *)
            break ;;
        esac
    done
    if [ $# -lt 2 ]; then
        enroot::usage exec 1
    fi
    pid="$1"
    shift

    runtime::exec "${pid}" \
      "$(IFS=$'\n'; echo ${environ[*]+"${environ[*]}"})" \
      "$@"
}

enroot::batch() {
    local conf=

    while [ $# -gt 0 ]; do
        case "$1" in
        -h|--help)
            enroot::usage batch 0 ;;
        --)
            shift; break ;;
        -?*)
            enroot::usage batch 1 ;;
        *)
            break ;;
        esac
    done
    if [ $# -lt 1 ]; then
        enroot::usage batch 1
    fi
    conf="$1"
    shift

    # Check for #ENROOT directives in the container configuration.
    if [ -n "${conf}" ] && [ -f "${conf}" ]; then
        common::checkcmd sed
        config::init <(sed -n '/^#[[:space:]]*ENROOT_/s/#//p' "${conf}") true; config::fini
    fi
    if [ -z "${ENROOT_ROOTFS-}" ]; then
        common::err "Missing directive ENROOT_ROOTFS in ${conf}"
    fi

    runtime::start "${ENROOT_ROOTFS}" "" "${conf}" "" "" "$@"
}

enroot::list() {
    local fancy=

    while [ $# -gt 0 ]; do
        case "$1" in
        -f|--fancy)
            fancy=y
            shift
            ;;
        -h|--help)
            enroot::usage list 0 ;;
        --)
            shift; break ;;
        -?*)
            enroot::usage list 1 ;;
        *)
            break ;;
        esac
    done

    runtime::list "${fancy}"
}

enroot::remove() {
    local name=

    while [ $# -gt 0 ]; do
        case "$1" in
        -f|--force)
            export ENROOT_FORCE_OVERRIDE=y
            shift
            ;;
        -h|--help)
            enroot::usage remove 0 ;;
        --)
            shift; break ;;
        -?*)
            enroot::usage remove 1 ;;
        *)
            break ;;
        esac
    done
    if [ $# -lt 1 ]; then
        enroot::usage remove 1
    fi

    for name in "$@"; do
        runtime::remove "${name}"
    done
}

enroot::bundle() {
    local image= filename= target= desc=

    while [ $# -gt 0 ]; do
        case "$1" in
        -a|--all)
            export ENROOT_BUNDLE_ALL=y
            shift
            ;;
        -c|--checksum)
            export ENROOT_BUNDLE_CHECKSUM=y
            shift
            ;;
        -f|--force)
            export ENROOT_FORCE_OVERRIDE=y
            shift
            ;;
        -o|--output)
            [ -z "${2-}" ] && enroot::usage bundle 1
            filename="$2"
            shift 2
            ;;
        --output=*)
            [ -z "${1#*=}" ] && enroot::usage bundle 1
            filename="${1#*=}"
            shift
            ;;
        -t|--target)
            [ -z "${2-}" ] && enroot::usage bundle 1
            target="$2"
            shift 2
            ;;
        --target=*)
            [ -z "${1#*=}" ] && enroot::usage bundle 1
            target="${1#*=}"
            shift
            ;;
        -d|--desc)
            [ -z "${2-}" ] && enroot::usage bundle 1
            desc="$2"
            shift 2
            ;;
        --desc=*)
            [ -z "${1#*=}" ] && enroot::usage bundle 1
            desc="${1#*=}"
            shift
            ;;
        -h|--help)
            enroot::usage bundle 0 ;;
        --)
            shift; break ;;
        -?*)
            enroot::usage bundle 1 ;;
        *)
            break ;;
        esac
    done
    if [ $# -ne 1 ]; then
        enroot::usage bundle 1
    fi
    image="$1"

    runtime::bundle "${image}" "${filename}" "${target}" "${desc}"
}

if [ $# -lt 1 ]; then
    enroot::usage help 1
fi
command="$1"; shift

case "${command}" in
version)
    enroot::version "$@" ;;
info)
    enroot::info "$@" ;;
digest)
    enroot::digest "$@" ;;
import)
    enroot::import "$@" ;;
load)
    enroot::load "$@" ;;
export)
    enroot::export "$@" ;;
create)
    enroot::create "$@" ;;
start)
    enroot::start "$@" ;;
exec)
    enroot::exec "$@" ;;
batch)
    enroot::batch "$@" ;;
list)
    enroot::list "$@" ;;
remove)
    enroot::remove "$@" ;;
bundle)
    enroot::bundle "$@" ;;
help)
    enroot::usage help 0 ;;
*)
    enroot::usage help 1 ;;
esac

exit 0
