#!/bin/sh

# Copyright (c) 2023 The PaSh Authors.
#
# Usage of this source code is governed by the MIT license, you can find the
# LICENSE file in the root directory of this project.
#
# https://github.com/binpash/try

TRY_VERSION="0.2.0"
TRY_COMMAND="${0##*/}"
export TRY_COMMAND

# exit status invariants
#
# 0 -- command ran
# 1 -- consistency error/failure
# 2 -- input error

################################################################################
# Run a command (in `$@`) in an overlay (in `$SANDBOX_DIR`)
################################################################################

try() {
    START_DIR="$PWD"

    if ! command -v findmnt >/dev/null
    then
        error "findmnt not found, please install util-linux" "$TRY_COMMAND" 2
    fi

    if [ "$SANDBOX_DIR" ]
    then
        ## If the name of a sandbox is given then we need to exit prematurely if its directory doesn't exist
        ! [ -d "$SANDBOX_DIR" ] && { error "could not find sandbox directory $SANDBOX_DIR" 2; }
    else
        ## Create a new sandbox if one was not given
        SANDBOX_DIR=$(mktemp -d)
    fi

    ## If the sandbox is not valid we exit early
    if ! sandbox_valid_or_empty "$SANDBOX_DIR"
    then
        error "given sandbox '$SANDBOX_DIR' is invalid" 1
    fi

    ## Make any directories that don't already exist, this is OK to do here
    ## because we have already checked if it valid.
    export SANDBOX_DIR

    # If we're in a docker container, we want to mount tmpfs on sandbox_dir, #136
    tmpdirname=$(mktemp -d)
    # tail -n +2 to ignore the first line with the column name
    tmpfstype=$(df --output=fstype "$tmpdirname" | tail -n +2)
    if [ "$tmpfstype" = "overlay" ] && [ "$(id -u)" -eq "0" ]
    then
        mount -t tmpfs tmpfs "$SANDBOX_DIR"
    fi

    mkdir -p "$SANDBOX_DIR/upperdir" "$SANDBOX_DIR/workdir" "$SANDBOX_DIR/temproot"

    ## Find all the directories and mounts that need to be mounted
    DIRS_AND_MOUNTS="$(mktemp)"
    export DIRS_AND_MOUNTS
    find / -maxdepth 1 >"$DIRS_AND_MOUNTS"
    findmnt --real -r -o target -n >>"$DIRS_AND_MOUNTS"
    sort -u -o "$DIRS_AND_MOUNTS" "$DIRS_AND_MOUNTS"

    # Calculate UPDATED_DIRS_AND_MOUNTS that contains the merge arguments in LOWER_DIRS
    UPDATED_DIRS_AND_MOUNTS="$(mktemp)"
    export UPDATED_DIRS_AND_MOUNTS
    while IFS="" read -r mountpoint
    do
        new_mountpoint=""
        OLDIFS=$IFS
        IFS=":"

        for lower_dir in $LOWER_DIRS
        do
            temp_mountpoint="$lower_dir/upperdir$mountpoint"
            if [ -n "$new_mountpoint" ]
            then
                # If new_mountpoint is not empty, append : and the temp_mountpoint
                new_mountpoint="$new_mountpoint:$temp_mountpoint"
            else
                # If new_mountpoint is empty, just set it to temp_mountpoint
                new_mountpoint="$temp_mountpoint"
            fi
        done
        IFS=$OLDIFS
        # Add the original mountpoint at the end
        new_mountpoint="${new_mountpoint:+$new_mountpoint:}$mountpoint"
        echo "$new_mountpoint" >> "$UPDATED_DIRS_AND_MOUNTS"
    done <"$DIRS_AND_MOUNTS"


    # we will overlay-mount each root directory separately (instead of all at once) because some directories cannot be overlayed
    # so we set up the mount points now
    #
    # KK 2023-06-29 This approach (of mounting each root directory separately) was necessary because we could not mount `/` in an overlay.
    #               However, this might be solvable using mergerfs/unionfs, allowing us to mount an overlay on a unionfs of the `/` once.
    #
    # findmnt
    # --real: only list real filesystems
    # -n: no header
    # -r: raw output
    # -o target: only print the mount target
    # then we want to exclude the root partition "/"
    while IFS="" read -r mountpoint
    do
        ## Only make the directory if the original is a directory too
        if [ -d "$mountpoint" ] && ! [ -L "$mountpoint" ]
        then
            # shellcheck disable=SC2174 # warning acknowledged, "When used with -p, -m only applies to the deepest directory."
            mkdir -m "$(stat -c %a "$mountpoint")" -p "${SANDBOX_DIR}/upperdir/${mountpoint}" "${SANDBOX_DIR}/workdir/${mountpoint}" "${SANDBOX_DIR}/temproot/${mountpoint}"
        fi
    done <"$DIRS_AND_MOUNTS"

    chmod "$(stat -c %a /)" "$SANDBOX_DIR/temproot"

    mount_and_execute="$(mktemp)"
    chroot_executable="$(mktemp)"
    try_mount_log="$(mktemp)"
    script_to_execute="$(mktemp)"

    export chroot_executable
    export try_mount_log
    export script_to_execute

    cat >"$mount_and_execute" <<"EOF"
#!/bin/sh

TRY_COMMAND="$TRY_COMMAND($0)"

## A wrapper of `mount -t overlay` to have cleaner looking code
make_overlay() {
    sandbox_dir="$1"
    lowerdirs="$2"
    overlay_mountpoint="$3"
    mount -t overlay overlay -o userxattr -o "lowerdir=$lowerdirs,upperdir=$sandbox_dir/upperdir/$overlay_mountpoint,workdir=$sandbox_dir/workdir/$overlay_mountpoint" "$sandbox_dir/temproot/$overlay_mountpoint"
}


devices_to_mount="tty null zero full random urandom"

## Mounts and unmounts a few select devices instead of the whole `/dev`
mount_devices() {
    sandbox_dir="$1"
    for dev in $devices_to_mount
    do
        touch "$sandbox_dir/temproot/dev/$dev"
        mount -o bind /dev/$dev "$sandbox_dir/temproot/dev/$dev"
    done
}

unmount_devices() {
    sandbox_dir="$1"
    for dev in $devices_to_mount
    do
        umount "$sandbox_dir/temproot/dev/$dev" 2>>"$try_mount_log"
        rm -f "$sandbox_dir/temproot/dev/$dev"
    done
}

## Try to autodetect union helper: {mergerfs | unionfs}
## Returns an empty string if no union helper is found
autodetect_union_helper() {
    if command -v mergerfs >/dev/null; then
        UNION_HELPER=mergerfs
    elif command -v unionfs >/dev/null; then
        UNION_HELPER=unionfs
    fi
}

# Detect if union_helper is set, if not, we try to autodetect them
if [ -z "$UNION_HELPER" ]
then
    ## Try to detect the union_helper (the variable could still be empty afterwards).
    autodetect_union_helper
fi

# actually mount the overlays
for mountpoint in $(cat "$UPDATED_DIRS_AND_MOUNTS")
do
    pure_mountpoint=${mountpoint##*:}

    ## We are not interested in mounts that are not directories
    if ! [ -d "$pure_mountpoint" ]
    then
        continue
    fi

    ## Symlinks
    if [ -L "$mountpoint" ]
    then
        ln -s $(readlink "$mountpoint") "$SANDBOX_DIR/temproot/$mountpoint"
        continue
    fi

    ## Don't do anything for the root and skip if it is /dev or /proc, we will mount it later
    case "$pure_mountpoint" in
        (/|/dev|/proc) continue;;
    esac

    # Try mounting everything normally
    make_overlay "$SANDBOX_DIR" "$mountpoint" "$pure_mountpoint" 2>>"$try_mount_log"
    # If mounting everything normally fails, we try using either using mergerfs or unionfs to mount them.
    if [ "$?" -ne 0 ]
    then
        ## If the overlay failed, it means that there is a nested mount inside the target mount, e.g., both `/home` and `/home/user/mnt` are mounts.
        ## To address this, we use unionfs/mergerfs (they support the same functionality) to show all mounts under the target mount as normal directories.
        ## Then we can normally make the overlay on the new union directory.
        ##
        ## KK 2023-06-29 Since this uses findmnt, it performs the union+overlay for both the outside and the inside mount.
        ##               In the best case scenario this is only causing extra work (the internal mount is already shown through the unionfs),
        ##                 but in the worst case this could lead to bugs due to the extra complexity (e.g., because we are doing mounts on top of each other).
        ##               We should try to investigate either:
        ##               1. Not doing another overlay if we have done it for a parent directory (we can keep around a list of overlays and skip if we are in a child)
        ##               2. Do one unionfs+overlay at the root `/` once and be done with it!

        if [ -z "$UNION_HELPER" ]
        then
            ## We can ignore this mountpoint, if the user program tries to use it, it will crash, but if not we can run normally
            printf "%s: Warning: Failed mounting $mountpoint as an overlay and mergerfs or unionfs not set and could not be found, see \"$try_mount_log\"\n" "$TRY_COMMAND" >&2
        else
            merger_dir=$(mktemp -d)

            ## Create a union directory
            "$UNION_HELPER" $mountpoint $merger_dir 2>>"$try_mount_log" ||
                printf "%s: Warning: Failed mounting $mountpoint via $UNION_HELPER, see \"$try_mount_log\"\n" "$TRY_COMMAND" >&2
            make_overlay "$SANDBOX_DIR" "$merger_dir" "$pure_mountpoint" 2>>"$try_mount_log" ||
            printf "%s: Warning: Failed mounting $mountpoint as an overlay via $UNION_HELPER, see \"$try_mount_log\"\n" "$TRY_COMMAND" >&2
        fi
    fi
done

## Mount a few select devices in /dev
mount_devices "$SANDBOX_DIR"

## Check if chroot_executable exists, #29
if ! [ -f "$SANDBOX_DIR/temproot/$chroot_executable" ]
then
    cp $chroot_executable "$SANDBOX_DIR/temproot/$chroot_executable"
fi

unshare --root="$SANDBOX_DIR/temproot" /bin/sh "$chroot_executable"
exitcode="$?"

# unmount the devices
sync
unmount_devices "$SANDBOX_DIR"

exit $exitcode
EOF

    # NB we substitute in the heredoc, so the early unsets are okay!
    cat >"$chroot_executable" <<EOF
#!/bin/sh

unset START_DIR SANDBOX_DIR UNION_HELPER DIRS_AND_MOUNTS TRY_EXIT_STATUS
unset script_to_execute chroot_executable try_mount_log

mount -t proc proc /proc &&
cd "$START_DIR" &&
. "$script_to_execute"
EOF

    echo "$@" >"$script_to_execute"

    # `$script_to_execute` need not be +x to be sourced
    chmod +x "$mount_and_execute" "$chroot_executable"

    # enable job control so interactive commands will play nicely with try asking for user input later(for committing). #5
    [ -t 0 ] && set -m

    # --mount: mounting and unmounting filesystems will not affect the rest of the system outside the unshare
    # --map-root-user: map to the superuser UID and GID in the newly created user namespace.
    # --user: the process will have a distinct set of UIDs, GIDs and capabilities.
    # --pid: create a new process namespace (needed fr procfs to work right)
    # --fork: necessary if we do --pid
    #         "Creation of a persistent PID namespace will fail if the --fork option is not also specified."
    # shellcheck disable=SC2086 # we want field splitting!
    unshare --mount --map-root-user --user --pid --fork $EXTRA_NS "$mount_and_execute"
    TRY_EXIT_STATUS=$?

    # remove symlink
    # first set temproot to be writible, rhel derivatives defaults / to r-xr-xr-x
    chmod 755 "${SANDBOX_DIR}/temproot"
    while IFS="" read -r mountpoint
    do
        if [  -L "$mountpoint" ]
        then
            rm "${SANDBOX_DIR}/temproot/${mountpoint}"
        fi
    done <"$DIRS_AND_MOUNTS"

    ################################################################################
    # commit?

    case "$NO_COMMIT" in
        (quiet)       ;;
        (show)        echo "$SANDBOX_DIR";;
        (commit)      commit;;
        (interactive) summary >&2
                      # shellcheck disable=SC2181
                      if [ "$?" -eq 0 ]
                      then
                          printf "\nCommit these changes? [y/N] " >&2
                          read -r DO_COMMIT
                          case "$DO_COMMIT" in
                              (y|Y|yes|YES) commit;;
                              (*)           echo "Not committing." >&2
                                            echo "$SANDBOX_DIR";;
                          esac
                      fi;;
    esac
}

################################################################################
# Summarize the overlay in `$SANDBOX_DIR`
################################################################################

summary() {
    if ! [ -d "$SANDBOX_DIR" ]
    then
        error "could not find directory $SANDBOX_DIR" 2
    elif ! [ -d "$SANDBOX_DIR/upperdir" ]
    then
        error "could not find directory $SANDBOX_DIR/upperdir" 1
    fi

    ## Finds all potential changes
    changed_files=$(find_upperdir_changes "$SANDBOX_DIR" "$IGNORE_FILE")
    summary_output=$(process_changes "$SANDBOX_DIR" "$changed_files")

    if [ -z "$summary_output" ]
    then
        return 1
    fi

    echo
    echo "Changes detected in the following files:"
    echo

    while IFS= read -r summary_line; do
        local_file="$(echo "$summary_line" | cut -c 4-)"
        case "$summary_line" in
            (ln*) echo "$local_file (symlink)";;
            (rd*) echo "$local_file (replaced with dir)";;
            (md*) echo "$local_file (created dir)";;
            (de*) echo "$local_file (deleted)";;
            (mo*) echo "$local_file (modified)";;
            (ad*) echo "$local_file (added)";;
        esac
    done <<EOF
$summary_output
EOF

    TRY_EXIT_STATUS=0
}

################################################################################
# Commit the results of an overlay in `$SANDBOX_DIR`
################################################################################

commit() {
    if ! [ -d "$SANDBOX_DIR" ]
    then
        error "could not find directory $SANDBOX_DIR" "$TRY_COMMAND" 2
    elif ! [ -d "$SANDBOX_DIR/upperdir" ]
    then
        error "could not find directory $SANDBOX_DIR/upperdir" 1
    fi

    changed_files=$(find_upperdir_changes "$SANDBOX_DIR" "$IGNORE_FILE")
    summary_output=$(process_changes "$SANDBOX_DIR" "$changed_files")

    TRY_EXIT_STATUS=0
    while IFS= read -r summary_line; do
        local_file="$(echo "$summary_line" | cut -c 4-)"
        changed_file="$SANDBOX_DIR/upperdir$local_file"
        case $summary_line in
            (ln*) ln -s "$(readlink "$changed_file")" "$local_file";;
            (rd*) rm -f "$local_file"; mkdir "$local_file";;
            (md*) mkdir "$local_file";;
            (de*) rm "$local_file";;
            (mo*) cp -fa "$changed_file" "$local_file";;
            (ad*) cp -fa "$changed_file" "$local_file";;
        esac

        # shellcheck disable=SC2181
        if [ "$?" -ne 0 ]
        then
            warn "couldn't commit $changed_file"
            TRY_EXIT_STATUS=1
        fi
    done <<EOF
$summary_output
EOF
}

################################################################################
# Checks if a file is an overlayfs whiteout file
################################################################################

is_whiteout_file() {
    file="$1"

    [ -c "$file" ] && ! [ -s "$file" ] && [ "$(stat -c %t,%T "$file")" = "0,0" ]
}

################################################################################
## Defines which changes we want to ignore in the summary and commit
################################################################################

ignore_changes() {
    ignore_file="$1"

    grep -v -f "$ignore_file"
}

################################################################################
## Lists all upperdir changes in raw format
################################################################################

find_upperdir_changes() {
    sandbox_dir="$1"
    ignore_file="$2"

    find "$sandbox_dir/upperdir/" -type f -o \( -type c -size 0 \) -o -type d -o -type l | ignore_changes "$ignore_file"
}

################################################################################
# Processes upperdir changes to an internal format that can be processed by summary and commit
#
# Output format:
#
#   XX PATH
#
#   where:
#     XX is a two character code for the modification
#     - rd: Replaced with a directory
#     - md: Created a directory
#     - de: Deleted a file
#     - mo: Modified a file
#     - ad: Added a file
#
#     PATH is the local/host path (i.e., without the upper
################################################################################

process_changes() {
    sandbox_dir="$1"
    changed_files="$2"

    while IFS= read -r changed_file
    do
        local_file="${changed_file#"$sandbox_dir/upperdir"}"
        if [ -L "$changed_file" ]
        then
            echo "ln $local_file"
        elif [ -d "$changed_file" ] && ! [ -d "${local_file}" ]
        then # new directory
            ## If something exists there, we need to delete it first
            if [ -e "$local_file" ]
            then
                echo "rd $local_file"
            else
                echo "md $local_file"
            fi
        elif is_whiteout_file "$changed_file"
        then # whiteout file
            echo "de $local_file"
        elif [ -f "$changed_file" ]
        then # normal file
            if [ -e "$local_file" ]
            then
                echo "mo $local_file"
            else
                echo "ad $local_file"
            fi
        fi
    done <<EOF
$changed_files
EOF
}

################################################################################
# Returns 0 if a sandbox is empty (fresh for use) or pre-existing and well formed
################################################################################

sandbox_valid_or_empty() {
    sandbox_dir="$1"

    if ! [ -d "$sandbox_dir/upperdir" ] && ! [ -d "$sandbox_dir/workdir" ] && ! [ -d "$sandbox_dir/temproot" ]
    then
        # No sandbox directory exists so we can happily return
        return 0
    fi

    # The sandbox already exists so we now need to check if it is valid
    # Validity requirements:
    # - no file exists in the temproot tree, i.e., all directories are empty
    #
    # TODO: Make this validity assertion tighter
    # KK 2023-06-28 workdir seems to be non-empty after a single use, is that expected?
    if [ "$(find "$sandbox_dir/temproot" -depth -not -type d)" ]
    then
        return 1
    fi

    return 0
}

################################################################################
# Emit a warning
################################################################################

warn() {
    msg="$1"

    printf "%s: %s\n" "$TRY_COMMAND" "$msg" >&2
}

################################################################################
# Emit a warning and exit
################################################################################

error() {
    msg="$1"
    exit_status="$2"

    warn "$msg"
    exit "$exit_status"
}

################################################################################
# Argument parsing
################################################################################

usage() {
    cat >&2 <<EOF
Usage: $TRY_COMMAND [-nvhyx] [-i PATTERN] [-D DIR] [-U PATH] [-L dir1:dir2:...] CMD [ARG ...]

  -n                don't commit or prompt for commit (overrides -y)
  -y                assume yes to all prompts (overrides -n)
  -x                prevent network access (by unsharing the network namespace)
  -i PATTERN        ignore paths that match PATTERN on summary and commit
  -D DIR            work in DIR (implies -n)
  -U PATH           path to unionfs helper (e.g., mergerfs, unionfs-fuse)
  -L dir1:dir2:...  specify multiple lower directories to merge (colon-separated, implies -n)
  -v                show version information (and exit)
  -h                show this usage message (and exit)

Subcommands:
  $TRY_COMMAND summary DIR   show the summary for the overlay in DIR
  $TRY_COMMAND commit DIR    commit the overlay in DIR
  $TRY_COMMAND explore DIR   start a shell inside the overlay in DIR
EOF
}

################################################################################
# Main entry point
################################################################################

# "interactive" - show nothing, interactively prompt on commit
# "show"        - show the resulting directory on stdout when we're done
# "quiet"       - do not show the result directory on stdout when we're done
# "commit"      - commit the result directory automatically when we're done
NO_COMMIT="interactive"

# Includes all patterns given using the `-i` flag; will be used with `grep -f`
IGNORE_FILE="$(mktemp)"

while getopts ":yvnhxi:D:U:L:" opt

do
    case "$opt" in
        (y)   NO_COMMIT="commit";;
        (n)   NO_COMMIT="show";;
        (i)   echo "$OPTARG" >>"$IGNORE_FILE";;
        (D)   if ! [ -d "$OPTARG" ]
              then
                  error "could not find sandbox directory '$OPTARG'" 2
              fi
              SANDBOX_DIR="$OPTARG"
              NO_COMMIT="quiet";;
        (L)   if [ -n "$LOWER_DIRS" ]
              then
                  error "the -L option has been specified multiple times" 2
              fi
              LOWER_DIRS="$OPTARG"
              NO_COMMIT="quiet";;
        (v)   echo "$TRY_COMMAND version $TRY_VERSION" >&2
              exit 0;;
        (U)   if ! [ -x "$OPTARG" ]
              then
                  error "could not find executable union helper '$OPTARG'" 2
              fi
              UNION_HELPER="$OPTARG"
              export UNION_HELPER;;
        (x)   EXTRA_NS="--net";;
        (h|*) usage
              exit 0;;
    esac
done

shift $((OPTIND - 1))

if [ "$#" -eq 0 ]
then
    usage
    exit 2
fi

TRY_EXIT_STATUS=1
case "$1" in
    (summary) : "${SANDBOX_DIR=$2}"
              summary;;
    (commit)  : "${SANDBOX_DIR=$2}"
              commit;;
    (explore) : "${SANDBOX_DIR=$2}"
              try "$SHELL";;
    (--)      shift
              try "$@";;
    (*)       try "$@";;
esac

exit "$TRY_EXIT_STATUS"