pipe-receive

#!/bin/bash

trap callstack ERR

set -o errexit \
    -o pipefail \
    -o nounset

global_msg_level=info
prog_name="${0##*/}"

function msg() {
    if [[ "$#" -ne 2 ]]; then
        printf "[%s] %s: (%s) %s" "$(logdate)" "$prog_name" "fatal" "incorrect number of arguments provided to msg function"
        exit 1
    fi

    local level="$1"
    local message="$2"
    local active_level="${global_msg_level:-none}"

    local -A levels=([none]=0 [fatal]=1 [error]=2 [warn]=3 [info]=4 [debug]=5)
    local num_level="${levels[$level]}"
    local num_active_level="${levels[$active_level]}"

    if [[ "$num_level" -le "${num_active_level}" ]]; then
        printf "[%s] %s: (%s) %s\n" "$(logdate)" "$prog_name" "$level" "$message" >&2
    fi
    if [[ "$level" == "fatal" ]]; then
        if [[ "${debug_mode:-false}" == "true" ]]; then
            callstack
        else
            exit 1
        fi
    fi
}

function logdate() {
    date +"%F %T"
}

function callstack() {
    echo --- Trace: ---
    local i
    # These arrays start at index 0, but 0 is this callstack function, so skip it.
    for ((i=1; i<${#FUNCNAME[*]}; i++)) do
        if [[ $i -ne 1 ]]; then
            echo -n "  "  
        fi
        echo "${BASH_SOURCE[$i]}: in \"${FUNCNAME[$i]}\" called from line ${BASH_LINENO[$i]}"
    done
    echo --------------
    exit 1
}

function must_have_cmds() {
    for cmd in "$@"; do
        command -v "$cmd" >/dev/null 2>/dev/null || msg fatal "this script needs a working copy of '$cmd' but could not find one"
    done
}

must_have_cmds tar date hostname


###########

# Arg Parsing -- not a lot

help_message="
usage: $prog_name [-h|-d] 

A script to quickly receive files from another user on the same node.

The user you are receiving the file from should run the equivalent give
 command, and should tell you the ID and password for the transfer.

This script will request the ID and password from you when you run it.

Transfer passwords are single-use, so don't worry about sharing them
 with someone else.

Note that this will not overwrite files with the same name you already
 have: if you have a file in the current directory called \"my_file\"
 and someone gives you a file called \"my_file\", this will be treated
 as an error.

Flags:
   -h, --help        Show this help.
   -d, --debug       Debug mode.
   -s, --skip        Skip the file listing and confirmation prompt.
"


function show_help_and_exit() {
    printf "%s" "$help_message"
    exit "${1:-0}"
}

canonical_args="$(
/usr/bin/getopt \
    -n pipe-receive \
    -l "help,debug,skip" \
    -o "hds" \
    -- \
    "$@"
)"

eval set -- "$canonical_args"

# Defaults
debug_mode="false"
skip_confirmation="false"

while true ; do
    case "$1" in
        -h|--help) show_help_and_exit 0; shift ;;
        -d|--debug) debug_mode="true"; shift ;;
        -s|--skip) skip_confirmation="true"; shift ;;
        --) shift ; break ;;
        *) msg fatal "invalid argument '$1'" ; exit 1 ;;
    esac
done

if [[ "$debug_mode" == "true" ]]; then
  # We don't have a proper debug mode, so just set the bash tracing options.
  set -v -x
else
  # If not in debug mode, remove the callstack printer from the error handler.
  trap ERR
fi


##########

# Okay, that's all the setup handled.



read -r -p "Please enter your transfer ID: " transfer_id 

transfer_fifo_dir="${TMPDIR:-/tmp}/pipe-transfer.$transfer_id"

if [[ ! -d "$transfer_fifo_dir" ]]; then
    this_hostname="$(hostname -s)"
    msg fatal "could not find a transfer with that id -- check sender node is $this_hostname"
fi

# Got ID, try password
result="incorrect password"

while [[ "$result" == "incorrect password" ]]; do
    read -r -p "Please enter transfer password: " transfer_password

    echo "$transfer_password" >"$transfer_fifo_dir/ctrl"

    result="$(cat "$transfer_fifo_dir/ctrl")"

    if [[ "$result" == "incorrect password" ]]; then
        msg warn "incorrect password"
    fi
done

if [[ "$result" != "correct password" ]]; then
    if [[ "$result" == "too many tries" ]]; then
        msg fatal "too many incorrect password attempts, quitting"
    else
        msg fatal "unknown status: $result"
    fi
fi

# So by this point we've either quit, or passed a correct password.

msg info "password correct, receiving identity and file count"
sending_user="$(cat -- "$transfer_fifo_dir/ctrl")"
file_count="$(cat -- "$transfer_fifo_dir/ctrl")"

if [[ "$sending_user" =~ ^sending\ user:[a-z0-9_-][a-z0-9_-]* ]]; then
    sending_user="${sending_user#sending user:}"
else
    msg fatal "invalid message received while trying to get sender identity"
fi

if [[ "$file_count" =~ ^sending\ file\ count:[0-9][0-9]* ]]; then
    file_count="${file_count#sending file count:}"
else
    msg fatal "invalid message received while trying to get file count"
fi


msg info "receiving list of files for approval"
file_list="$(cat -- "$transfer_fifo_dir/data")"
if [[ "$skip_confirmation" == "true" ]]; then
    msg info "skipping approval"
    confirm_option="y"
else
    printf "\n-- File List --\n%s\n---------------\n\n" "$file_list"
    if [[ "$file_count" -eq 1 ]]; then
        printf "User %s wants to copy this file to you.\n" "$sending_user"
    else
        printf "User %s wants to copy these %s files to you.\n" "$sending_user" "$file_count"
    fi
    read -s -r -N 1 -p "Would you like to accept? [press y for yes, anything else to cancel]" confirm_option
    printf "\n"
fi

if [[ "$confirm_option" == "y" ]]; then
    msg info "confirming with sender"
    echo "confirm" >"$transfer_fifo_dir/ctrl"
else
    echo "reject" >"$transfer_fifo_dir/ctrl"
    msg fatal "rejected listing, will quit"
fi

# Work out whether we have a working pv but if we don't, no worries.
if command -v pv >/dev/null 2>/dev/null; then
    PV=pv
elif command -v /shared/ucl/apps/pv/1.6.6/bin/pv >/dev/null 2>/dev/null; then
    PV=/shared/ucl/apps/pv/1.6.6/bin/pv
else
    PV="(none)"
    msg warn "could not find pv -- transfer speed display will not be shown"
fi

# Run the actual transfer, piping to tar from the data fifo.
msg info "receiving files..."
transfer_status=problem
if [[ "$PV" == "(none)" ]]; then
    tar -xvk <"$transfer_fifo_dir/data" && transfer_status=success
else
    "$PV" <"$transfer_fifo_dir/data" | tar -xvk && transfer_status=success
fi

# Report back so the sender can tell whether there was a problem.
if [[ "$transfer_status" == "problem" ]]; then
    echo "problem" > "$transfer_fifo_dir/ctrl"
    msg fatal "there was a problem transferring the files"
else
    echo "success" > "$transfer_fifo_dir/ctrl"
    msg info "transfer complete"
fi