#!/usr/bin/env bash

# Merge an overlay directory into a squashfs
#
# Usage: merge-squashfs input.squashfs output.squashfs overlay-dir
#

set -e

LC_ALL=C
LANG=C
export LC_ALL LANG

# Account for Mac differences
if [[ "$(uname)" = "Darwin" ]]; then
    # Use gstat on the Mac, but make sure that it exists first.
    STAT=$(command -v gstat || echo "/usr/local/bin/gstat")
    if [[ ! -e "$STAT" ]]; then
        echo "merge-squashfs: ERROR: Please install gstat first"
        echo
        echo "For example:"
        echo "   brew install coreutils"
        exit 1
    fi
else
    STAT=stat
fi

# This implementation might be slow
perm_to_num() {
	string=$1

	let perms=0

	[[ "${string}" = ?r???????? ]] && perms=$(( perms + 0400 ))
	[[ "${string}" = ??w??????? ]] && perms=$(( perms + 0200 ))
	[[ "${string}" = ???x?????? ]] && perms=$(( perms + 0100 ))
	[[ "${string}" = ???s?????? ]] && perms=$(( perms + 04100 ))
	[[ "${string}" = ???S?????? ]] && perms=$(( perms + 04000 ))
	[[ "${string}" = ????r????? ]] && perms=$(( perms + 040 ))
	[[ "${string}" = ?????w???? ]] && perms=$(( perms + 020 ))
	[[ "${string}" = ??????x??? ]] && perms=$(( perms + 010 ))
	[[ "${string}" = ??????s??? ]] && perms=$(( perms + 02010 ))
	[[ "${string}" = ??????S??? ]] && perms=$(( perms + 02000 ))
	[[ "${string}" = ???????r?? ]] && perms=$(( perms + 04 ))
	[[ "${string}" = ????????w? ]] && perms=$(( perms + 02 ))
	[[ "${string}" = ?????????x ]] && perms=$(( perms + 01 ))
	[[ "${string}" = ?????????t ]] && perms=$(( perms + 01001 ))
	[[ "${string}" = ?????????T ]] && perms=$(( perms + 01000 ))

    printf '%o' $perms
}

owner_to_uid_gid() {
    echo "$1" | sed -e "s/\// /g"
}

get_path() {
    # trim the squashfs-root and escape any spaces
    echo "$*" | sed -e 's/squashfs-root//' -e 's/ /\\ /g'
}

device_file() {
    perm=$1
    owner=$2

    # Depending on the major and minor numbers, you either get
    # major,minor or major, minor. Remove the comma and account
    # for the optional space.
    if [[ $3 =~ ,.+ ]]; then
        major_minor=$(echo "$3" | sed -e 's/,/ /')
        filename=$6
    else
        major_minor="$(echo "$3" | sed -e 's/,/ /') $4"
        filename=$7
    fi

    echo "$(get_path "$filename") ${perm:0:1} $(perm_to_num "$perm") $(owner_to_uid_gid "$owner") $major_minor"
}

symlink_file() {
    # symlinks are the handled the same as files currently.
    echo "$(get_path "$6") m $(perm_to_num "$1") $(owner_to_uid_gid "$2")"
}

directory_file() {
    pathname=$(get_path "$6")
    # mksquashfs doesn't support setting permissions and ownership
    # on the root directory
    if [[ ! -z $pathname ]]; then
        echo "$pathname m $(perm_to_num "$1") $(owner_to_uid_gid "$2")"
    fi
}

regular_file() {
    echo "$(get_path "$6") m $(perm_to_num "$1") $(owner_to_uid_gid "$2")"
}

unsquash_to_pseudofile() {
    input=$1
    while read -r perms owner rest; do
        case $perms in
            c*)
                device_file "$perms" "$owner" $rest
                ;;
            b*)
                device_file "$perms" "$owner" $rest
                ;;
            l*)
                symlink_file "$perms" "$owner" $rest
                ;;
            d*)
                directory_file "$perms" "$owner" $rest
                ;;
            -*)
                regular_file "$perms" "$owner" $rest
                ;;
            *)
                # Ignore
                ;;
        esac
    done < <(unsquashfs -n -ll "$input")
}

find_to_pseudofile() {
    inputdir=$1
    while read -r mode path; do
        path=$(echo "$path" | sed -e "s^$inputdir^^" -e "s/ /\\\\ /g")
        if [[ ! -z $path ]]; then
            echo "$path m $mode 0 0"
        fi
    done < <(find "$inputdir" -exec $STAT -c "%a %n" "{}" ";")
}


# "readlink -f" implementation for BSD
# This code was extracted from the Elixir shell scripts
readlink_f () {
    cd "$(dirname "$1")" > /dev/null
    filename="$(basename "$1")"
    if [[ -h "$filename" ]]; then
        readlink_f "$(readlink "$filename")"
    else
        echo "$(pwd -P)/$filename"
    fi
}

if [[ $# -ne 3 ]]; then
    echo "Usage: merge-squashfs input.squashfs output.squashfs overlay-dir"
    exit 1
fi

input_squashfs=$(readlink_f "$1")
output_squashfs=$(readlink_f "$2")
overlay_dir=$(readlink_f "$3")

if [[ ! -f $input_squashfs ]]; then
    echo "Can't open $input_squashfs"
    exit 1
fi
if [[ ! -d $overlay_dir ]]; then
    echo "Can't find $overlay_dir"
    exit 1
fi

workdir=$(mktemp -d 2>/dev/null || mktemp -d -t 'merge-squashfs-tmp')
pushd "$workdir" >/dev/null

unsquash_to_pseudofile "$input_squashfs" >pseudofile.in
find_to_pseudofile "$overlay_dir" >>pseudofile.in

# remove repeated entries to avoid warnings from mksquashfs
awk '!x[$1]++' pseudofile.in > pseudofile

unsquashfs "$input_squashfs" >/dev/null 2>/dev/null
cp -Rf "$overlay_dir/." "$workdir/squashfs-root"
mksquashfs squashfs-root "$output_squashfs" -pf pseudofile -noappend -no-recovery -no-progress

# cleanup
popd >/dev/null
rm -fr "$workdir"