keichi/pi-oven

View on GitHub
oven

Summary

Maintainability
Test Coverage
#!/bin/bash

# Copyright (c) 2016-2021 Keichi Takahashi <keichi.t@me.com>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

set -u

# Meta information
readonly VERSION=1.0.0
readonly PROGNAME=$(basename "$0")

# Some ANSI color codes
readonly C_RED='\033[0;31m'
readonly C_GREEN='\033[0;32m'
readonly C_YELLOW='\033[0;33m'
readonly C_START='\033[0;32m'
readonly C_RESET='\033[m'

# Global variables
opt_resize=
opt_disk_image=
opt_script=
opt_interactive=false
# Path to loop mount the disk image
readonly opt_mount_point=/mnt/oven
# Partition number of the boot partition
opt_boot_part_num=1
# Partition number of the root partition
opt_root_part_num=2

# exit_with_usage - Print usage and exit
function exit_with_usage() {
    cat << EOS >&2
Usage: ${PROGNAME} [options] src [dst]
  ${PROGNAME}: Customize RaspberryPi disk images

  Options:
  -r, --resize SIZE    Resize the root file system before customization
  -s, --script PATH    Path to a shell script for customization
  -i, --interactive    Start an interactive shell for customization
  --bootpart [1-4]     Partition number of the boot partition (default=1)
  --rootpart [1-4]     Partition number of the root partition (default=2)
  --version            Print version information
  -h, --help           Show usage
EOS
    exit 1
}

# exit_with_version -- Print version info and exit
function exit_with_version() {
    echo "$VERSION"
    exit 0
}

# log loglevel msgbody - Output log
function log() {
    local loglevel=$1
    local msgbody=$2

    case $loglevel in
        INFO)  echo -en "[${C_START}${C_GREEN}${loglevel}${C_RESET}]";;
        WARN)  echo -en "[${C_START}${C_YELLOW}${loglevel}${C_RESET}]";;
        ERROR) echo -en "[${C_START}${C_RED}${loglevel}${C_RESET}]";;
        *)     echo -en "[${loglevel}]";;
    esac
    echo " ${msgbody}"
}

# check_image image -- Check if image exists
function check_image() {
    local image=$1
    if ! [[ -s $image ]] || ! [[ -r $image ]]; then
        log "ERROR" "Disk image ${image} is not readable"
        exit 1
    fi
}

# check_deps - Check if dependent tools exist
function check_deps() {
    local failed=false

    # Check if fallocate, parted and kpartx exist in PATH
    for name in fallocate parted kpartx; do
        if ! type "$name" &> /dev/null; then
            log "ERROR" "${name} is missing"
            failed=true
        fi
    done

    # Some dependencies are missing
    if $failed; then
        log "ERROR" "Cannot continue; please install dependencies first"
        exit 1
    fi
}

# resize_root size image - Resize root partition
function resize_root() {
    local size_mb=$1
    local size_b=$((size_mb * 1000 * 1000))
    local image=$2

    log "INFO" "Resizing disk image file ${image} to ${size_mb}MB"

    log "INFO" "Extending disk image file"
    fallocate -l "$size_b" "$image"

    log "INFO" "Extending root partition"
    if ! parted -h | grep resizepart > /dev/null; then
        log "WARN" "parted does not support resizepart subcommand"
        log "WARN" "Falling back to rm and mkpart"

        local regexp="s/^\s*$opt_root_part_num\s\+\([0-9]\+\).*/\1/p"
        local offset
        offset=$(parted -s "$image" unit s p | sed -n "$regexp")
        parted -s "$image" rm "$opt_root_part_num" mkpart primary ext4 "${offset}s" 100%
    else
        parted -s "$image" resizepart "$opt_root_part_num" "$size_mb"
    fi
}

# is_mounted - Check if disk image is mounted
function is_mounted() {
    if mount | grep "$opt_mount_point" > /dev/null; then
        return 0
    else
        return 1
    fi
}

# mount_image image - Loop mount root partition of disk image
function mount_image() {
    local image=$1

    if is_mounted; then
        return 0
    fi

    # Create device mappers
    log "INFO" "Creating device mapper for ${image}"
    local kx_out
    # Make sure device mappers are deleted
    kpartx -dsv "$image" > /dev/null
    kx_out=$(kpartx -asv "$image")

    # Extract device mapper path
    local regex="^add map ([[:alnum:]]+)"
    if [[ $(echo "$kx_out" | sed -n ${opt_boot_part_num}p) =~ $regex ]]; then
        local dm_boot=/dev/mapper/${BASH_REMATCH[1]}
    else
        log "ERROR" "Could not find device mapper for boot partition"
        exit 1
    fi
    if [[ $(echo "$kx_out" | sed -n ${opt_root_part_num}p) =~ $regex ]]; then
        local dm_root=/dev/mapper/${BASH_REMATCH[1]}
    else
        log "ERROR" "Could not find device mapper for root partition"
        exit 1
    fi

    # Resize root filesystem if needed
    e2fsck -fy "$dm_root" &> /dev/null
    resize2fs "$dm_root" &> /dev/null

    # Make sure mount point exists
    if ! [[ -d $opt_mount_point  ]]; then
        mkdir -p "$opt_mount_point"
    fi

    # Loop mount root partition
    log "INFO" "Mounting ${dm_root} to ${opt_mount_point}"
    if ! mount "$dm_root" "$opt_mount_point" &> /dev/null; then
        log "ERROR" "Failed to mount ${dm_root}"
        exit 1
    fi

    # Loop mount boot partition
    log "INFO" "Mounting ${dm_boot} to ${opt_mount_point}/boot"
    if ! mount "$dm_boot" "$opt_mount_point/boot" &> /dev/null; then
        log "ERROR" "Failed to mount ${dm_boot}"
        exit 1
    fi

    # Mount devfs, sysfs and procfs
    log "INFO" "Mounting devfs, sysfs and procfs"
    mount --bind /dev "${opt_mount_point}/dev"
    mount --bind /sys "${opt_mount_point}/sys"
    mount -t proc none "${opt_mount_point}/proc"
}

# umount_image image - Unmount disk image
function umount_image() {
    local image=$1

    if ! is_mounted; then
        return 0
    fi

    # Unmount devfs, sysfs and procfs
    log "INFO" "Unmouting devfs, sysfs and procfs"
    umount "${opt_mount_point}/dev"
    umount "${opt_mount_point}/sys"
    umount "${opt_mount_point}/proc"

    log "INFO" "Unmounting ${opt_mount_point}/boot"
    umount "$opt_mount_point/boot"

    log "INFO" "Unmounting ${opt_mount_point}"
    umount "$opt_mount_point"

    log "INFO" "Deleting device mapper for ${image}"
    kpartx -dsv "$image" > /dev/null
}

# check_privileges - Check if we have root privilege
function check_privileges() {
    if ! [ "$(whoami)" == "root" ]; then
        log "ERROR" "${PROGNAME} requires root privilege"
        exit 1
    fi
}

# prepare_chroot - Preparation before running chroot
function prepare_chroot() {
    log "INFO" "Preparing chroot environment"
}

# cleanup_chroot - Cleanup after running chroot
function cleanup_chroot() {
    log "INFO" "Cleaning up chroot environment"
}

# provision_interactive - Run interactive provisioning
function provision_interactive() {
    prepare_chroot

    log "INFO" "Interactive provisioning: launching interactive shell"
    chroot $opt_mount_point /bin/bash

    cleanup_chroot
}

# provision_script - Run shell script provisioning
function provision_script() {
    local script=$1

    if ! [[ -s $script ]] || ! [[ -r $script ]]; then
        log "ERROR" "Provisioning script ${script} is not readable"
        exit 1
    fi

    prepare_chroot
    cp "$script" "${opt_mount_point}/usr/bin/provision.sh"

    log "INFO" "Script provisioning: executing provisioning script ${script}"
    if chroot $opt_mount_point /bin/bash -ex /usr/bin/provision.sh; then
        log "INFO" "Successfully executed provisioning script"
    else
        log "ERROR" "Provisioning script returned error"
    fi

    rm -f "${opt_mount_point}/usr/bin/provision.sh"
    cleanup_chroot
}

# parse_args *args - Parse command line arguments
function parse_args() {
    options=$(getopt -o r:s:ih --long resize:,script:,interactive,rootpart:,bootpart:,version,help -- "$@")
    if [[ $? != 0 ]]; then
        exit_with_usage
    fi
    eval set -- "$options"

    while true; do
        case "$1" in
            -r | --resize)      readonly opt_resize=$2; shift 2;;
            -s | --script)      readonly opt_script=$(pwd)/"$2"; shift 2;;
            -i | --interactive) readonly opt_interactive=true; shift;;
            --bootpart)         readonly opt_boot_part_num=$2; shift 2;;
            --rootpart)         readonly opt_root_part_num=$2; shift 2;;
            -h | --help)        exit_with_usage; break;;
            --version)          exit_with_version; break;;
            --)                 shift; break;;
            *)                  exit_with_usage; break;;
        esac
    done

    if ! $opt_interactive && [[ $opt_script == ""  ]]; then
        log "ERROR" "Either --script or --interactive must be set"
        exit_with_usage
    fi

    if (( opt_boot_part_num < 1 )) || (( opt_boot_part_num > 4 )); then
        log "ERROR" "--bootpart must be a number between 1 to 4"
        exit_with_usage
    fi

    if (( opt_root_part_num < 1 )) || (( opt_root_part_num > 4 )); then
        log "ERROR" "--rootpart must be a number between 1 to 4"
        exit_with_usage
    fi

    if [[ $# -eq 1 ]]; then
        readonly opt_disk_image=$(pwd)/"$1";
    elif [[ $# -eq 2 ]]; then
        log "INFO" "Copying disk image file $1 to $2"
        cp "$1" "$2"
        readonly opt_disk_image=$(pwd)/"$2";
    else
        exit_with_usage
    fi
}

# main *args - Main function
function main() {
    parse_args "$@"

    trap 'umount_image ${opt_disk_image}' 0

    check_privileges
    check_deps
    check_image "$opt_disk_image"

    if [[ $opt_resize != "" ]]; then
        resize_root "$opt_resize" "$opt_disk_image"
    fi

    mount_image "$opt_disk_image"

    if [[ $opt_script != "" ]]; then
        provision_script "$opt_script"
    fi

    if $opt_interactive; then
        provision_interactive
    fi

    umount_image "$opt_disk_image"

    exit 0
}

main "$@"