#!/usr/bin/bash
#
# SPDX-License-Identifier: LGPL-2.1+
#
# Create application containers from OCI images

set -eu

# Make sure the usual locations are in PATH
export PATH="$PATH:/usr/sbin:/usr/bin:/sbin:/bin"

# Check for required binaries
for bin in skopeo umoci jq; do
  if ! command -v $bin >/dev/null 2>&1; then
    echo "ERROR: Missing required tool: $bin" 1>&2
    exit 1
  fi
done

LOCALSTATEDIR=/var
LXC_TEMPLATE_CONFIG=/usr/share/lxc/config
LXC_HOOK_DIR=/usr/share/lxc/hooks
MOUNT_HELPER=""
MOUNTED_WORKDIR=""

# Some useful functions
cleanup() {
  if [ -d "${LXC_ROOTFS}.tmp" ]; then
    rm -Rf "${LXC_ROOTFS}.tmp"
  fi
  if [ -n "${MOUNTED_WORKDIR}" ]; then
    echo "${MOUNT_HELPER} unmount ${MOUNTED_WORKDIR}" >&2
    "${MOUNT_HELPER}" umount "${MOUNTED_WORKDIR}"
    MOUNTED_WORKDIR=""
  fi
}

in_userns() {
  [ -e /proc/self/uid_map ] || { echo no; return; }
  while read -r line; do
    fields="$(echo "$line" | awk '{ print $1 " " $2 " " $3 }')"
    if [ "${fields}" = "0 0 4294967295" ]; then
      echo no;
      return;
    fi
    if echo "${fields}" | grep -q " 0 1$"; then
      echo userns-root;
      return;
    fi
  done < /proc/self/uid_map

  if [ -e /proc/1/uid_map ]; then
    if [ "$(cat /proc/self/uid_map)" = "$(cat /proc/1/uid_map)" ]; then
      echo userns-root
      return
    fi
  fi
  echo yes
}

getconfigpath() {
  local basedir="$1" mfpath="$2" cdigest=""
  mtdigest=$(jq -c -r '.config.mediaType' < "$mfpath")
  if [ "$mtdigest" = "application/vnd.oci.empty.v1+json" ]; then
    echo ""
    return 0
  fi

  # Ok we have the image config digest, now get the config ref from the manifest.
  # shellcheck disable=SC2039
  cdigest=$(jq -c -r '.config.digest' < "$mfpath")
  if [ -z "${cdigest}" ]; then
    echo "container config not found" >&2
    return
  fi

  # cdigest is '<hashtype>:<hash>', so 'ht' gets type, hv gets value.
  local ht="${cdigest%%:*}" hv="${cdigest#*:}" p=""
  p="$basedir/blobs/$ht/$hv"
  if [ ! -f "$p" ]; then
    echo "config file did not exist for digest $cdigest" >&2
    return 1
  fi
  echo "$p"
}

getmanifestpath() {
  local basedir="$1" ref="$2" p=""
  # if given 'sha256:<hash>' then return the blobs/sha256/hash
  case "$ref" in
    sha256:*)
      p="$basedir/blobs/sha256/${ref#sha256:}"
      [ -f "$p" ] && echo "$p" && return 0
      echo "could not find manifest path to blob $ref. file did not exist: $p" >&2
      return 1
      ;;
  esac
  # find the reference by annotation
  local blobref="" hashtype="" hashval=""
  blobref=$(jq -c -r --arg q "$ref" '.manifests[] | if .annotations."org.opencontainers.image.ref.name" == $q then .digest else empty end' < "${basedir}/index.json")
  # blobref is 'hashtype:hash'
  hashtype="${blobref%%:*}"
  hashval="${blobref#*:}"
  p="$basedir/blobs/$hashtype/$hashval"
  [ -f "$p" ] && echo "$p" && return 0
  echo "did not find manifest for $ref. file did not exist: $p" >&2
  return 1
}

getlayermediatype() {
  jq -c -r '.layers[0].mediaType' <"$1"
}

# Get entrypoint from oci image. Use sh if unspecified
getep() {
  if [ "$#" -eq 0 ]; then
    echo "/bin/sh"
    return
  fi

  configpath="$1"
  if [ "${configpath}" = "" ]; then
    echo "/bin/sh"
    return
  fi

  ep=$(jq -c '.config.Entrypoint[]?'< "${configpath}" | tr '\n' ' ')
  cmd=$(jq -c '.config.Cmd[]?'< "${configpath}" | tr '\n' ' ')
  if [ -z "${ep}" ]; then
    ep="${cmd}"
    if [ -z "${ep}" ]; then
      ep="/bin/sh"
    fi
  elif [ -n "${cmd}" ]; then
    ep="${ep} ${cmd}"
  fi

  echo "${ep}"
  return
}

# get environment from oci image.
getenv() {
  if [ "$#" -eq 0 ]; then
    return
  fi

  configpath="$1"
  if [ "${configpath}" = "" ]; then
    return
  fi

  env=$(jq -c -r '.config.Env[]'< "${configpath}")

  echo "${env}"
  return
}

# check var is decimal.
isdecimal() {
  var="$1"
  if [ "${var}" -eq "${var}" ] 2> /dev/null; then
    return 0
  else 
    return 1
  fi
}

# get uid, gid from oci image.
getuidgid() {
  configpath="$1"
  rootpath="$2"
  passwdpath="${rootpath}/etc/passwd"
  grouppath="${rootpath}/etc/group"

  if [ "${configpath}" = "" ]; then
    user=0
    group=0
    echo "${user:-0} ${group:-0}"
    return
  fi
  usergroup=$(jq -c -r '.config.User' < "${configpath}")
  # shellcheck disable=SC2039
  usergroup=(${usergroup//:/ })

  user=${usergroup[0]:-0}
  if ! isdecimal "${user}"; then
    if [ -f ${passwdpath} ]; then
      user=$(grep "^${user}:" "${passwdpath}" | awk -F: '{print $3}')
    else
      user=0
    fi
  fi

  group=${usergroup[1]:-}
  if [ -z "${group}" ]; then
    if [ -f "${passwdpath}" ]; then
      group=$(grep "^[^:]*:[^:]*:${user}:" "${passwdpath}" | awk -F: '{print $4}')
    else
      group=0
    fi
  elif ! isdecimal "${group}"; then
    if [ -f "${grouppath}" ]; then
      group=$(grep "^${group}:" "${grouppath}" | awk -F: '{print $3}')
    else
      group=0
    fi
  fi

  echo "${user:-0} ${group:-0}"
  return
}

# get cwd from oci image.
getcwd() {
  if [ "$#" -eq 0 ]; then
    echo "/"
    return
  fi

  configpath="$1"
  if [ "${configpath}" = "" ]; then
    echo "/"
    return
  fi

  cwd=$(jq -c -r '.config.WorkingDir // "/"' < "${configpath}")

  echo "${cwd}"
  return
}

is_same_mountpoint() {
  local path1="$1"
  local path2="$2"
  local mount1
  local mount2

  # make sure bind mounts are covered
  mount1=$(df --output=target "$path1" | tail -n 1)
  mount2=$(df --output=target "$path2" | tail -n 1)
  [ "$mount1" = "$mount2" ]
}

usage() {
  cat <<EOF
LXC container template for OCI images

Special arguments:
[ -h | --help ]: Print this help message and exit.

Required arguments:
[ -u | --url <url> ]: The OCI image URL

Optional arguments:
[ --username <username> ]: The username for the registry
[ --password <password> ]: The password for the registry
[ --mount-helper <command> ]: program that will be used to mount. default will be detected from mediatype

     mount-helper is expected to support being called with 'mount'
     and 'umount' subcommands as below:

        mount-helper mount --persist <upperdir> <oci_dir>:<oci_name> <mountpoint>
        mount-helper umount <mountpoint>

     The --persist <upperdir> flag tells the mount helper to create a writable overlay, with a read-only
     filesystem as lowerdir and <upperdir> as upperdir, where <upperdir> is a filesystem path

LXC internal arguments (do not pass manually!):
[ --name <name> ]: The container name
[ --path <path> ]: The path to the container
[ --rootfs <rootfs> ]: The path to the container's rootfs
[ --mapped-uid <map> ]: A uid map (user namespaces)
[ --mapped-gid <map> ]: A gid map (user namespaces)
EOF
  return 0
}

if ! options=$(getopt -o u:h -l help,url:,username:,password:,no-cache,dhcp,name:,path:,rootfs:,mapped-uid:,mapped-gid:,mount-helper: -- "$@"); then
    usage
    exit 1
fi
eval set -- "$options"

OCI_URL=""
OCI_USERNAME=
OCI_PASSWORD=
OCI_USE_CACHE="true"
OCI_USE_DHCP="false"

LXC_MAPPED_GID=
LXC_MAPPED_UID=
LXC_NAME=
LXC_PATH=
LXC_ROOTFS=

while :; do
  case "$1" in
    -h|--help)    usage && exit 1;;
    -u|--url)     OCI_URL=$2; shift 2;;
    --username)   OCI_USERNAME=$2; shift 2;;
    --password)   OCI_PASSWORD=$2; shift 2;;
    --no-cache)   OCI_USE_CACHE="false"; shift 1;;
    --dhcp)       OCI_USE_DHCP="true"; shift 1;;
    --name)       LXC_NAME=$2; shift 2;;
    --path)       LXC_PATH=$2; shift 2;;
    --rootfs)     LXC_ROOTFS=$2; shift 2;;
    --mapped-uid) LXC_MAPPED_UID=$2; shift 2;;
    --mapped-gid) LXC_MAPPED_GID=$2; shift 2;;
    --mount-helper) MOUNT_HELPER=$2; shift 2;;
    *)            break;;
  esac
done

# Check that we have all variables we need
if [ -z "$LXC_NAME" ] || [ -z "$LXC_PATH" ] || [ -z "$LXC_ROOTFS" ]; then
  echo "ERROR: Not running through LXC" 1>&2
  exit 1
fi

if [ -z "$OCI_URL" ]; then
  echo "ERROR: no OCI URL given"
  exit 1
fi

if [ -n "$OCI_PASSWORD" ] && [ -z "$OCI_USERNAME" ]; then
  echo "ERROR: password given but no username specified"
  exit 1
fi

if [ "${OCI_USE_CACHE}" = "true" ]; then
  if ! skopeo copy --help | grep -q 'dest-shared-blob-dir'; then
    echo "INFO: skopeo doesn't support blob caching"
    OCI_USE_CACHE="false"
  fi
fi

USERNS=$(in_userns)

if [ "$USERNS" = "yes" ]; then
  if [ -z "$LXC_MAPPED_UID" ] || [ "$LXC_MAPPED_UID" = "-1" ]; then
    echo "ERROR: In a user namespace without a map" 1>&2
    exit 1
  fi
fi

OCI_DIR="$LXC_PATH/oci"
if [ "${OCI_USE_CACHE}" = "true" ]; then
  if [ "$USERNS" = "yes" ]; then
    DOWNLOAD_BASE="${HOME}/.cache/lxc"
  else
    DOWNLOAD_BASE="${LOCALSTATEDIR}/cache/lxc"
  fi
else
  DOWNLOAD_BASE="$OCI_DIR"
fi

if [ "${OCI_USE_CACHE}" = "true" ]; then
  if ! is_same_mountpoint "$LXC_PATH" "$DOWNLOAD_BASE"; then
    OCI_USE_CACHE="false"
    DOWNLOAD_BASE="$OCI_DIR"
  fi
fi

mkdir -p "${DOWNLOAD_BASE}"

# Trap all exit signals
trap cleanup EXIT HUP INT TERM

# Download the image
# shellcheck disable=SC2039
skopeo_args=("--remove-signatures" "--insecure-policy")
if [ -n "$OCI_USERNAME" ]; then
  CREDENTIALS="${OCI_USERNAME}"

  if [ -n "$OCI_PASSWORD" ]; then
    CREDENTIALS="${CREDENTIALS}:${OCI_PASSWORD}"
  fi

  # shellcheck disable=SC2039
  skopeo_args+=(--src-creds "${CREDENTIALS}")
fi

OCI_NAME="$LXC_NAME"
if [ "${OCI_USE_CACHE}" = "true" ]; then
  skopeo_args+=(--dest-shared-blob-dir "${DOWNLOAD_BASE}")
  mkdir -p "${OCI_DIR}/blobs/"
  ln -s "${DOWNLOAD_BASE}/sha256" "${OCI_DIR}/blobs/sha256"
fi

skopeo copy "${skopeo_args[@]}" "${OCI_URL}" "oci:${OCI_DIR}:${OCI_NAME}"

mfpath=$(getmanifestpath "${OCI_DIR}" "${OCI_NAME}")
OCI_CONF_FILE=$(getconfigpath "${OCI_DIR}" "$mfpath")
mediatype=$(getlayermediatype "$mfpath")
echo "mfpath=$mfpath conf=$OCI_CONF_FILE" 1>&2
echo "mediatype=$mediatype" >&2

case "$mediatype" in
  application/vnd.*.image.layer.squashfs*)
    MOUNT_HELPER="atomfs"
    ;;
  application/vnd.puzzlefs.image.rootfs.*)
    MOUNT_HELPER="puzzlefs"
    ;;
esac

case "$mediatype" in
  #application/vnd.oci.image.layer.v1.tar+gzip
  application/vnd.oci.image.layer.v1.tar*)
    echo "Unpacking tar rootfs" 2>&1
    # shellcheck disable=SC2039
    umoci_args=("")
    if [ -n "$LXC_MAPPED_UID" ] && [ "$LXC_MAPPED_UID" != "-1" ]; then
      # shellcheck disable=SC2039
      umoci_args+=(--rootless)
    fi
    # shellcheck disable=SC2039
    # shellcheck disable=SC2068
    umoci --log=error unpack ${umoci_args[@]} --image "${OCI_DIR}:${OCI_NAME}" "${LXC_ROOTFS}.tmp"
    find "${LXC_ROOTFS}.tmp/rootfs" -mindepth 1 -maxdepth 1 -exec mv '{}' "${LXC_ROOTFS}/" \;
    ;;
  #application/vnd.stacker.image.layer.squashfs+zstd+verity
  application/vnd.*.image.layer.squashfs*|application/vnd.puzzlefs.image.rootfs.*)
    if [ -z "${MOUNT_HELPER}" ]; then
      echo "MOUNT_HELPER not detected for $mediatype"
      exit 1
    fi
    if ! command -v "${MOUNT_HELPER}" >/dev/null 2>&1; then
      echo "media type $mediatype requires $MOUNT_HELPER" >&2
      exit 1
    fi
    MOUNT_HELPER_UPPERDIR="$LXC_PATH/upper"
    echo "$MOUNT_HELPER mount --persist ${MOUNT_HELPER_UPPERDIR} ${OCI_DIR}:${OCI_NAME} $LXC_ROOTFS" >&2
    "$MOUNT_HELPER" mount --persist "${MOUNT_HELPER_UPPERDIR}" "${OCI_DIR}:${OCI_NAME}" "$LXC_ROOTFS"
    MOUNTED_WORKDIR="$LXC_ROOTFS"
    ;;
  *)
    echo "Unknown media type $mediatype" >&2
    exit 1
    ;;
esac

LXC_CONF_FILE="${LXC_PATH}/config"
entrypoint=$(getep "${OCI_CONF_FILE}")
echo "lxc.execute.cmd = '${entrypoint}'" >> "${LXC_CONF_FILE}"
echo "lxc.mount.auto = proc:mixed sys:mixed cgroup:mixed" >> "${LXC_CONF_FILE}"

case "$mediatype" in
  application/vnd.*.image.layer.squashfs*|application/vnd.puzzlefs.image.rootfs.*)
    echo "lxc.hook.version = 1" >> "${LXC_CONF_FILE}"
    # shellcheck disable=SC2016
    echo "lxc.hook.pre-mount = $MOUNT_HELPER mount --persist ${MOUNT_HELPER_UPPERDIR}" \
        '${LXC_ROOTFS_PATH}/../oci:${LXC_NAME} ${LXC_ROOTFS_PATH}' \
        >> "${LXC_CONF_FILE}";;
esac

environment=$(getenv "${OCI_CONF_FILE}")
# shellcheck disable=SC2039
while read -r line; do
  echo "lxc.environment = ${line}" >> "${LXC_CONF_FILE}"
done <<< "${environment}"

if [ -e "${LXC_TEMPLATE_CONFIG}/common.conf" ]; then
  echo "lxc.include = ${LXC_TEMPLATE_CONFIG}/common.conf" >> "${LXC_CONF_FILE}"
fi

if [ -n "$LXC_MAPPED_UID" ] && [ "$LXC_MAPPED_UID" != "-1" ] && [ -e "${LXC_TEMPLATE_CONFIG}/userns.conf" ]; then
  echo "lxc.include = ${LXC_TEMPLATE_CONFIG}/userns.conf" >> "${LXC_CONF_FILE}"
fi

if [ -e "${LXC_TEMPLATE_CONFIG}/oci.common.conf" ]; then
  echo "lxc.include = ${LXC_TEMPLATE_CONFIG}/oci.common.conf" >> "${LXC_CONF_FILE}"
fi

if [ "${OCI_USE_DHCP}" = "true" ]; then
  echo "lxc.hook.start-host = ${LXC_HOOK_DIR}/dhclient" >> "${LXC_CONF_FILE}"
  echo "lxc.hook.stop = ${LXC_HOOK_DIR}/dhclient" >> "${LXC_CONF_FILE}"
fi

echo "lxc.uts.name = ${LXC_NAME}" >> "${LXC_CONF_FILE}"
# set the hostname
cat <<EOF > "${LXC_ROOTFS}/etc/hostname"
${LXC_NAME}
EOF

# set minimal hosts
cat <<EOF > "${LXC_ROOTFS}/etc/hosts"
127.0.0.1   localhost
127.0.1.1   ${LXC_NAME}
::1     localhost ip6-localhost ip6-loopback
fe00::0 ip6-localnet
ff00::0 ip6-mcastprefix
ff02::1 ip6-allnodes
ff02::2 ip6-allrouters
EOF

# shellcheck disable=SC2039
uidgid=($(getuidgid "${OCI_CONF_FILE}" "${LXC_ROOTFS}" ))
# shellcheck disable=SC2039
echo "lxc.init.uid = ${uidgid[0]}" >> "${LXC_CONF_FILE}"
# shellcheck disable=SC2039
echo "lxc.init.gid = ${uidgid[1]}" >> "${LXC_CONF_FILE}"

cwd=$(getcwd "${OCI_CONF_FILE}")
echo "lxc.init.cwd = ${cwd}" >> "${LXC_CONF_FILE}"

if [ -n "$LXC_MAPPED_UID" ] && [ "$LXC_MAPPED_UID" != "-1" ]; then
  chown "$LXC_MAPPED_UID" "$LXC_PATH/config" "$LXC_PATH/fstab" > /dev/null 2>&1 || true
fi
if [ -n "$LXC_MAPPED_GID" ] && [ "$LXC_MAPPED_GID" != "-1" ]; then
  chgrp "$LXC_MAPPED_GID" "$LXC_PATH/config" "$LXC_PATH/fstab" > /dev/null 2>&1 || true
fi

exit 0