#!/bin/bash

#
# This script takes care of remapping the UIDs and GIDs of paths of non-static
# users and groups when moving a filesystem from one firmware install to a new
# one, where the UIDs and GIDs of non-static users and groups may be different.
#
# The script is resistent to interruptions due to power cycles or uncontrolled
# reboots, it will just restart and complete its job (although some operations
# may be done twice).
#
# To prepare the filesystem for remapping the following needs to be done while
# running on the "old" firmware install.
#
#  - Check if the $SENTINEL_NAME file exists at the root of the filesystem,
#    if it does then a previous remapping is still pending so migration to
#    the new firmware install must be aborted until the previous remapping is
#    completed.
#  - Copy the /etc/passwd and /etc/group files to $USER_DB_NAME and
#    $GROUP_DB_NAME at the root of the filesystem.
#  - Create an empty file $SENTINEL_NAME at the root of the filesystem
#
# Any errors during the remapping process are cumulatively logged to the
# $ERRORLOG_NAME file at the root of the filesystem. If the file does not
# exist then it means no remapping errors ever occured.
#
# See below for the actual values of the above variables.
#

# If non-empty run verbosely
VERBOSE=

# If non-empty to not make any changes
DRYRUN=

# Some constants
SENTINEL_NAME=.remap-ids
DONE_NAME=.remap-ids.done
ERRORLOG_NAME=.remap-ids.errors
USER_DB_NAME=.remap-ids.passwd
GROUP_DB_NAME=.remap-ids.group
PATH_LIST_NAME=.remap-ids.path-list

# Error out on any unexpected error
set -e
# Catch all errors in pipelines
set -o pipefail

function unexpected_error() {
    echo "ERROR: terminated due to unexpected error" >&2
}
trap unexpected_error ERR

# Generate the list of paths which potentially need to be remapped, along with
# their UID and GID, in the following format: UID colon GID colon path NUL
# arg1: the mountpoint of the filesystem to remap
# arg2: the path to the path listing
function generate_file_list() {
    # files which have a user and group with common users and groups which have static UIDs are omitted
    find "$1" -xdev \( \( \! -user root -a \! -user www-data \) -o \( \! -group root -a \! -group www-data \) \) -fprintf "$2" '%U:%G:%p\0'
}

# Returns the unique UID:GID pairs from the path listing, one per line
# arg1: the path to the path listing
function get_id_pairs() {
    cut -z -d: -f1-2 < "$1" | sort -zu | tr '\0' '\n'
}

# Sets the new owner for the paths in the path listing matching the given UID:GID pair
# arg1: UID:GID pair (e.g., "920:933")
# arg2: new owner as user:group
# arg3: the path to the path listing
function chown_matching_paths() {
    local id_pair="$1"
    local new_owner="$2"
    local path_list="$3"

    grep -za "^${id_pair}:" "$path_list" | cut -z -d: -f3- | xargs -0r ${VERBOSE:+-t} ${DRYRUN:+true} chown --no-dereference "$new_owner" --
}

# Does the actual ID remapping
# arg1: the filesystem root
function do_remap() {
    local fsroot="$1"
    local old_user_db="$fsroot"/$USER_DB_NAME
    local old_group_db="$fsroot"/$GROUP_DB_NAME
    local path_list="$fsroot"/$PATH_LIST_NAME
    local id_pairs p old_uid old_gid user group new_uid new_gid new_owner
    local ret=0

    if ! id_pairs="$(get_id_pairs "$path_list")"; then
        echo "ERROR: failed getting ID pairs for path listing" >&2
        return 1
    fi
    for p in $id_pairs; do
        old_uid="${p%:*}"
        old_gid="${p##*:}"
        user="$(awk -F: -v key="$old_uid" '{ if ($3 == key) print $1 }' "$old_user_db")"
        group="$(awk -F: -v key="$old_gid" '{ if ($3 == key) print $1 }' "$old_group_db")"
        if [ -z "$user" ]; then
            echo "ERROR: failed finding user for old UID '$old_uid'" >&2
            new_uid=
            ret=1
        else
            new_uid="$(getent passwd "$user" | cut -d: -f3)"
            if [ -z "$new_uid" ]; then
                echo "ERROR: failed finding new UID for user '$user'" >&2
                ret=1
            fi
        fi
        if [ -z "$group" ]; then
            echo "ERROR: failed finding group for old GID '$old_gid'" >&2
            ret=1
            new_gid=
        else
            new_gid="$(getent group "$group" | cut -d: -f3)"
            if [ -z "$new_gid" ]; then
                echo "ERROR: failed finding new GID for group '$group'" >&2
                ret=1
            fi
        fi
        if [ -n "$new_uid" ] && [ $new_uid -ne $old_uid ]; then
            : # need to remap at least UID
        elif [ -n "$new_gid" ] && [ $new_gid -ne $old_gid ]; then
            : # need to remap GID
        else
            [ -z "$VERBOSE" ] || echo "no need to remap or nothing to remap for '${old_uid}:${old_gid}' to '${user}:${group}' / '${new_uid}:${new_gid}'"
            continue # nothing to remap
        fi
        [ -z "$VERBOSE" ] || echo "remapping '${old_uid}:${old_gid}' to '${user}:${group}' / '${new_uid}:${new_gid}'"
        # NOTE: with chown the owner "user:" means user and its login group, to change just the user the colon is to be omitted
        if [ -n "$new_uid" ]; then
            [ -n "$new_gid" ] && new_owner="${user}:${group}" || new_owner="${user}"
        else
            new_owner=":${group}"
        fi
        if ! chown_matching_paths "${old_uid}:${old_gid}" "$new_owner" "$path_list"; then
            echo "ERROR: failed changing owner of files with UID:GID '${old_uid}:${old_gid}' to user:group '${user}:${group}'" >&2
            ret=1
        fi
    done

    # Make sure all the changes to the filesystem make it to stable storage before we return
    sync -f "$fsroot" || ret=1

    return $ret
}

# Does the initial setup. It is done only once before the process starts,
# if the process is interrupted this will do nothing.
# arg1: the filesystem root
function do_initialize() {
    local fsroot="$1"
    local sentinel="$fsroot"/$SENTINEL_NAME
    local old_user_db="$fsroot"/$USER_DB_NAME
    local old_group_db="$fsroot"/$GROUP_DB_NAME
    local path_list="$fsroot"/$PATH_LIST_NAME

    if ! [ -d "$fsroot" ]; then
        echo "ERROR: fs root '$fsroot' on which to remap IDs is not a directory" >&2
        return 1
    fi
    if ! [ -s "$old_user_db" ]; then
        echo "ERROR: old user DB '$old_user_db' is empty or does not exist" >&2
        return 1
    fi
    if ! [ -s "$old_group_db" ]; then
        echo "ERROR: old group DB '$old_group_db' is empty or does not exist" >&2
        return 1
    fi

    # make sure the sentinel exists and is submitted to storage
    touch "$sentinel" && sync "$fsroot"

    [ -s "$fsroot"/$ERRORLOG_NAME ] && \
        echo "WARNING: restarting remapping, but there were remapping errors previously" >> "$fsroot"/$ERRORLOG_NAME && \
        sync "$fsroot"/$ERRORLOG_NAME

    if [ -f "$path_list" ]; then
        [ -s "$fsroot"/$ERRORLOG_NAME ] && echo "INFO: reusing exising path list '$path_list'" >&2
        echo "INFO: reusing exising path list '$path_list'"
        return
    fi

    if ! generate_file_list "$fsroot" "$path_list".tmp || ! sync "$path_list".tmp; then
        echo "ERROR: failed generating the path listing" >&2
        rm -f "$path_list".tmp
        return 1
    fi

    if ! mv "$path_list".tmp "$path_list" || ! sync "$fsroot"; then
        echo "ERROR: failed comitting the path listing" >&2
        return 1
    fi
}

# Removes the data files and sentinel in the proper order
# arg1: the filesystem root
function do_cleanup() {
    local fsroot="$1"

    # The path listing must not be re-created after some remapping has been done, otherwise
    # the final UIDs and GIDs will be wrong, so cleanup must ensure the file cannot be re-created.
    touch "$fsroot"/$DONE_NAME
    sync "$fsroot"
    rm -f "$fsroot"/$USER_DB_NAME "$FSROOT"/$GROUP_DB_NAME
    sync "$fsroot"
    rm -f "$fsroot"/$PATH_LIST_NAME
    [ -s "$fsroot"/$ERRORLOG_NAME ] || rm -f "$fsroot"/$ERRORLOG_NAME
    [ -f "$fsroot"/$ERRORLOG_NAME ] && sync "$fsroot"/$ERRORLOG_NAME
    sync "$fsroot"
    rm -f "$fsroot"/$SENTINEL_NAME
    sync "$fsroot"
    rm -f "$fsroot"/$DONE_NAME
    sync "$fsroot"
}

function show_usage() {
    echo "$0: [options] fs-root"
    echo
    echo "Does remapping of UIDs on GIDs of all paths under fs-root which have changed"
    echo "UID or GID between the old user or group DB and the current user and group IDs."
    echo "The $USER_DB_NAME and $GROUP_DB_NAME files must exist under fs-root, with the"
    echo "old user and group DBs, respectively."
    echo
    echo "Options"
    echo -e "  -h\tShow this help message and exit"
    echo -e "  -v\tRun verbosely"
    echo -e "  -n\tDry-run, do not make any modifications"
}

while getopts :vnh opt; do
    case "$opt" in
        v)
            VERBOSE=1
            ;;
        n)
            DRYRUN=1
            ;;
        h)
            show_usage "$0";
            exit 0
            ;;
        ':')
            echo "ERROR: missing argument for option -${OPTARG}" >&2
            exit 1
            ;;
        '?')
            echo "ERROR: unknown option -${OPTARG}" >&2
            exit 1
            ;;
        *)
            echo "ERROR: unexpected error while parsing options" >&2
            exit 1
    esac
done
shift $((OPTIND-1))

FSROOT="$1"
if [ -z "$FSROOT" -o $# -gt 1 ]; then
    show_usage "$0"
    exit 1
fi

if [ -f "$FSROOT"/$DONE_NAME ]; then
    echo "INFO: remapping was already done, cleaning up"
    do_cleanup "$FSROOT"
    exit 0
fi

if ! do_initialize "$FSROOT" 2>> "$FSROOT"/$ERRORLOG_NAME; then
    echo "ERROR: failed initialization, check $FSROOT/$ERRORLOG_NAME" >&2
    do_cleanup "$FSROOT"
    exit 1
fi
if ! do_remap "$FSROOT" 2>> "$FSROOT"/$ERRORLOG_NAME; then
    echo "ERROR: some errors encountered during remapping, check $FSROOT/$ERRORLOG_NAME" >&2
    do_cleanup "$FSROOT"
    exit 1
fi

do_cleanup "$FSROOT"
