#!/bin/bash

#
# template script for generating Arch Linux container for LXC
#

#
# lxc: linux Container library

# Authors:
# Alexander Vladimirov <alexander.idkfa.vladimirov@gmail.com>
# John Lane <lxc@jelmail.com>

# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.

# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.

# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA

# Detect use under userns (unsupported)

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

# defaults
arch=$(uname -m)
default_path="/var/lib/lxc"
default_locale="en-US.UTF-8"
default_timezone="UTC"
pacman_config="/etc/pacman.conf"
common_config="/usr/share/lxc/config/common.conf"
shared_config="/usr/share/lxc/config/archlinux.common.conf"

# by default, install 'base' except the kernel
pkg_blacklist="linux"
base_packages=()
for pkg in $(pacman -Sqg base); do
   [ "${pkg_blacklist#*$pkg}" = "$pkg_blacklist" ] && base_packages+=($pkg)
done
declare -a additional_packages

# split comma-separated string into an array
# ${1} - string to split
# ${2} - separator (default is ",")
# ${result} - result value on success
split_string() {
    local ifs=${IFS}
    IFS="${2:-,}"
    read -a result < <(echo "${1}")
    IFS=${ifs}
    return 0
}

[ -f /etc/arch-release ] && is_arch=true

# Arch-specific preconfiguration for container
configure_arch() {
    # on ArchLinux, read defaults from host systemd configuration
    if [ "${is_arch}" ]; then
        cp -p /etc/locale.conf /etc/locale.gen "${rootfs_path}/etc/"
    else
        echo "LANG=${default_locale}" > "${rootfs_path}/etc/locale.conf"
        if [ -e "${rootfs_path}/etc/locale.gen" ]; then
            sed -i 's@^#\(en_US\.UTF-8\)@\1@' "${rootfs_path}/etc/locale.gen"
            if [ ! "${default_locale}" = "en_US.UTF-8" ]; then
                echo "${default_locale} ${default_locale##*.}" >> \
                "${rootfs_path}/etc/locale.gen"
            fi
        fi
    fi

    # hostname and nameservers
    echo "${name}" > "${rootfs_path}/etc/hostname"
    while read r; do
       [ "${r#nameserver}" = "$r" ] || echo "$r"
    done < /etc/resolv.conf > "${rootfs_path}/etc/resolv.conf"

    # chroot and configure system
    arch-chroot "${rootfs_path}" /bin/bash -s << EOF
mkdir /run/lock
locale-gen
ln -s /usr/share/zoneinfo/${default_timezone} /etc/localtime
# set default boot target
ln -s /lib/systemd/system/multi-user.target /etc/systemd/system/default.target
# override getty@.service for container ttys
sed -e 's/^ConditionPathExists=/# ConditionPathExists=/' \
    -e 's/After=dev-%i.device/After=/' \
    < /lib/systemd/system/getty\@.service \
    > /etc/systemd/system/getty\@.service
# initialize pacman keyring
pacman-key --init
pacman-key --populate archlinux
EOF
    # enable getty on active ttys
    local nttys=$(cat "${config_path}/config" ${shared_config} ${common_config} | grep "^lxc.tty" | head -n1 | cut -d= -f2 | tr -d "[:blank:]")
    local devttydir=$(cat "${config_path}/config" ${shared_config} ${common_config} | grep "^lxc.devttydir" | head -n1 | cut -d= -f2 | tr -d "[:blank:]")
    local devtty=""
    # bind getty instances to /dev/<devttydir>/tty* if lxc.devttydir is set
    [ -n "${devttydir}" ] && devtty="${devttydir}-"
    if [ ${nttys:-0} -gt 1 ]; then
      ( cd "${rootfs_path}/etc/systemd/system/getty.target.wants"
        for i in $(seq 1 $nttys); do ln -sf "../getty@.service" "getty@${devtty}tty${i}.service"; done )
    fi
    # update securetty to allow console login if devttydir is set
    if [ -n "${devttydir}" ]; then
        for i in $(seq 1 ${nttys:-1}); do
            echo "${devttydir}/tty${i}" >> "${rootfs_path}/etc/securetty"
        done
    fi
    [ -n "${devttydir}" ] && echo "${devttydir}/console" >> "${rootfs_path}/etc/securetty"
    # Arch default configuration allows only tty1-6 for login
    [ ${nttys:-0} -gt 6 ] && echo \
      "You may want to modify container's /etc/securetty \
      file to allow root logins on tty7 and higher"
    return 0
}

# write container configuration files
copy_configuration() {
    mkdir -p "${config_path}"
    local config="${config_path}/config"
    echo "lxc.utsname = ${name}" >> "${config}"
    grep -q "^lxc.arch" "${config}" 2>/dev/null \
        || echo "lxc.arch = ${arch}" >> "${config}"
    grep -q "^lxc.rootfs" "${config}" 2>/dev/null \
        || echo "lxc.rootfs = ${rootfs_path}" >> "${config}"
    [ -e "${shared_config}" ] \
        && echo "lxc.include = ${shared_config}" >> "${config}"
    if [ $? -ne 0 ]; then
        echo "Failed to configure container"
        return 1
    fi
    return 0
}

# install packages within container chroot
install_arch() {
    [ "${arch}" != "$(uname -m)" ] && different_arch=true

    if [ "${different_arch}" = "true" ]; then
        container_pacman_config=$(mktemp)
        container_mirrorlist=$(mktemp)
        sed -e "s:Architecture =.*:Architecture = ${arch}:g"  \
            -e "s:/etc/pacman.d/mirrorlist:${container_mirrorlist}:g" \
            "${pacman_config}" > "${container_pacman_config}"
        sed -e "s:\(x86_64\|\$arch\):${arch}:g" \
            /etc/pacman.d/mirrorlist > "${container_mirrorlist}"

        pacman_config="${container_pacman_config}"
    fi

    if ! pacstrap -dcGC "${pacman_config}" "${rootfs_path}" \
            ${base_packages[@]}; then
        echo "Failed to install container packages"
        return 1
    fi

    if [ "${different_arch}" = "true" ]; then
        sed -i -e "s:Architecture =.*:Architecture = ${arch}:g" \
            "${rootfs_path}"/etc/pacman.conf
        cp "${container_mirrorlist}" "${rootfs_path}"/etc/pacman.d/mirrorlist
        rm "${container_pacman_config}" "${container_mirrorlist}"
    fi

    [ -d "${rootfs_path}/lib/modules" ] && ldconfig -r "${rootfs_path}"
    return 0
}

usage() {
    cat <<EOF
usage:
    ${1} -n|--name=<container_name> [-p|--path=<path>] [-a|--arch=<arch of the container>]
        [-r|--root_password=<root password>] [-P|--packages=<pkg1,pkg2,...>]
        [-e|--enable_units=unit1,unit2...] [-d|--disable_units=unit1,unit2...]
        [-c|--config=<pacman config path>] [-h|--help]

Mandatory args:
  -n,--name           container name, used to as an identifier for that container from now on
Optional args:
  -p,--path           path to where the container rootfs will be created (${default_path})
  --rootfs            path for actual container rootfs, (${default_path/rootfs)
  -P,--packages       preinstall additional packages, comma-separated list
  -e,--enable_units   enable systemd services, comma-separated list
  -d,--disable_units  disable systemd services, comma-separated list
  -c,--config         use specified pacman config when installing container packages
  -a,--arch           use specified architecture instead of host's architecture
  -r,--root_password  set container root password
  -h,--help           print this help
EOF
    return 0
}

options=$(getopt -o hp:P:e:d:n:c:a:r: -l help,rootfs:,path:,packages:,enable_units:,disable_units:,name:,config:,arch:,root_password:,mapped-uid:,mapped-gid: -- "${@}")
if [ ${?} -ne 0 ]; then
    usage $(basename ${0})
    exit 1
fi
eval set -- "${options}"

mapped_uid=-1
mapped_gid=-1
while true
do
    case "${1}" in
    -h|--help)          usage ${0} && exit 0;;
    -p|--path)          path=${2}; shift 2;;
    -n|--name)          name=${2}; shift 2;;
    --rootfs)           rootfs_path=${2}; shift 2;;
    -P|--packages)      additional_packages=${2}; shift 2;;
    -e|--enable_units)  enable_units=${2}; shift 2;;
    -d|--disable_units) disable_units=${2}; shift 2;;
    -c|--config)        pacman_config=${2}; shift 2;;
    -a|--arch)          arch=${2}; shift 2;;
    -r|--root_password) root_passwd=${2}; shift 2;;
    --mapped-uid)       mapped_uid=$2; shift 2;;
    --mapped-gid)       mapped_gid=$2; shift 2;;
    --)             shift 1; break ;;
    *)              break ;;
    esac
done

if [ -z "${name}" ]; then
    echo "missing required 'name' parameter"
    exit 1
fi

type pacman >/dev/null 2>&1
if [ ${?} -ne 0 ]; then
    echo "'pacman' command is missing, refer to wiki.archlinux.org for information about installing pacman"
    exit 1
fi

if [ -z "${path}" ]; then
    path="${default_path}/${name}"
fi

if [ "${EUID}" != "0" ]; then
    echo "This script should be run as 'root'"
    exit 1
fi

if [ -z "$rootfs_path" ]; then
    rootfs_path="${path}/rootfs"
fi
config_path="${path}"

revert() {
    echo "Interrupted, cleaning up"
    lxc-destroy -n "${name}"
    rm -rf "${path}/${name}"
    rm -rf "${default_path}/${name}"
    exit 1
}

trap revert SIGHUP SIGINT SIGTERM

copy_configuration
if [ ${?} -ne 0 ]; then
    echo "failed to write configuration file"
    rm -rf "${config_path}"
    exit 1
fi

if [ ${#additional_packages[@]} -gt 0 ]; then
    split_string ${additional_packages}
    base_packages+=(${result[@]})
fi

mkdir -p "${rootfs_path}"
install_arch
if [ ${?} -ne 0 ]; then
    echo "failed to install Arch Linux"
    rm -rf "${config_path}" "${path}"
    exit 1
fi

configure_arch
if [ ${?} -ne 0 ]; then
    echo "failed to configure Arch Linux for a container"
    rm -rf "${config_path}" "${path}"
    exit 1
fi

if [ ${#enable_units[@]} -gt 0 ]; then
    split_string ${enable_units}
    for unit in ${result[@]}; do
        [ "${unit##*.}" = "service" ] || unit="${unit}.service"
        ln -s "/usr/lib/systemd/system/${unit}" \
            "${rootfs_path}/etc/systemd/system/multi-user.target.wants/"
    done
fi

if [ ${#disable_units[@]} -gt 0 ]; then
    split_string ${disable_units}
    for unit in ${result[@]}; do
        [ "${unit##*.}" = "service" ] || unit="${unit}.service"
        ln -s /dev/null "${rootfs_path}/etc/systemd/system/${unit}"
    done
fi

if [ -n "${root_passwd}" ]; then
    echo "root:${root_passwd}" | chroot "${rootfs_path}" chpasswd
fi

if [ $mapped_uid -ne -1 ]; then
    chown $mapped_uid "$config_path"
    chown -R $mapped_uid "$rootfs_path"
fi
if [ $mapped_gid -ne -1 ]; then
    chgrp $mapped_gid    "$config_path"
    chgrp -R $mapped_gid "$rootfs_path"
fi

cat << EOF
Arch Linux container ${name} is successfully created! The configuration is
stored in ${config_path}/config. Please refer to https://wiki.archlinux.org for
information about configuring Arch Linux.
EOF