#!/bin/bash

# These can be overridden when running this script
set_default_variables()
{
    HOSTNAME=${HOSTNAME:-$(hostname)}
    BACKUP_HOST=${BACKUP_HOST:-backup.jim.sh}
    BACKUP_PORT=${BACKUP_PORT:-222}
    BACKUP_USER=${BACKUP_USER:-jim-backups}
    BACKUP_REPO=${BACKUP_REPO:-borg/${HOSTNAME}}
    SYSTEMD_UNIT=${SYSTEMD_UNIT:-borg-backup}

    # Use stable host ID in case MAC address changes.
    # Note that this host ID is only used to manage locks, so it's
    # not crucial that it remains stable.
    UUID=$(python3 -c 'import uuid;print(uuid.getnode())')
    HOSTID=${BORG_HOST_ID:-"${HOSTNAME}@${UUID}"}

    log "Configuration:"
    log " HOSTNAME           Local hostname: \033[1m${HOSTNAME}"
    log " HOSTID              Local host ID: \033[1m${HOSTID}"
    log " BACKUP_USER    Backup server user: \033[1m${BACKUP_USER}"
    log " BACKUP_HOST    Backup server host: \033[1m${BACKUP_HOST}"
    log " BACKUP_PORT    Backup server port: \033[1m${BACKUP_PORT}"
    log " BACKUP_REPO     Server repository: \033[1m${BACKUP_REPO}"
    log " SYSTEMD_UNIT    Systemd unit name: \033[1m${SYSTEMD_UNIT}"
    for i in $(seq 15 -1 1); do
        printf "\rPress ENTER or wait $i seconds to continue... \b"
        if read -t 1 ; then break; fi
    done
}

# Main dir is where this repo was checked out
BORG_DIR="$(realpath "$(dirname "$0")")"
cd "${BORG_DIR}"

BORG_BIN="${BORG_DIR}/bin/borg.$(uname -m)"

function error_handler() {
    echo "Error at $1 line $2:"
    echo -n '>>> ' ; tail -n +"$2" < "$1" | head -1
    echo "... exited with code $3"
    exit "$3"
}
trap 'error_handler ${BASH_SOURCE} ${LINENO} $?' ERR
set -o errexit
set -o errtrace

# Create pip environment
setup_venv()
{
    if ! which pipenv >/dev/null 2>&1 ; then
        error "pipenv not found, try: sudo apt install pipenv"
    fi
    mkdir -p .venv
    pipenv install
}

# Create shell script with environment variables
create_borg_vars()
{
    VARS=${BORG_DIR}/vars.sh

    # These variables are used elsewhere in this script
    BORG_REPO="ssh://${BACKUP_USER}@${BACKUP_HOST}:${BACKUP_PORT}/./${BACKUP_REPO}"
    BORG=${BORG_DIR}/borg.sh
    SSH=$BORG_DIR/ssh

    cat >"$VARS" <<EOF
export BACKUP_USER=${BACKUP_USER}
export BACKUP_HOST=${BACKUP_HOST}
export BACKUP_PORT=${BACKUP_PORT}
export BACKUP_REPO=${BACKUP_REPO}
export HOSTNAME=${HOSTNAME}
export BORG_REPO=${BORG_REPO}
export BORG_HOST_ID=${HOSTID}
export BORG_PASSCOMMAND="cat ${BORG_DIR}/passphrase"
export BORG_DIR=${BORG_DIR}
export SSH=${SSH}
export BORG=${BORG}
export BORG_BIN=${BORG_BIN}
export SYSTEMD_UNIT=${SYSTEMD_UNIT}
EOF
    if ! "$BORG" -h >/dev/null ; then
        error "Can't run the borg wrapper; does borg work?"
    fi
}

# Copy templated files, filling in templates as needed
install_templated_files()
{
    DOCS="README.md"
    SCRIPTS="notify.sh logs.sh"
    for i in ${DOCS} ${SCRIPTS}; do
        sed -e "s!\${HOSTNAME}!${HOSTNAME}!g" \
            -e "s!\${BORG_DIR}!${BORG_DIR}!g" \
            -e "s!\${BORG_BIN}!${BORG_BIN}!g" \
            -e "s!\${BACKUP_USER}!${BACKUP_USER}!g" \
            -e "s!\${BACKUP_HOST}!${BACKUP_HOST}!g" \
            -e "s!\${BACKUP_PORT}!${BACKUP_PORT}!g" \
            -e "s!\${BACKUP_REPO}!${BACKUP_REPO}!g" \
            -e "s!\${SYSTEMD_UNIT}!${SYSTEMD_UNIT}!g" \
            templates/$i > $i
    done
    chmod +x ${SCRIPTS}
}

# Update local paths in scripts
update_paths()
{
    sed -i\
        -e "1c#!${BORG_DIR}/.venv/bin/python" \
        backup.py
}

# See if we're just supposed to update an existing install, or recovering
parse_args()
{
    RECOVER=0
    UPDATE=0
    if [ "$1" == "--recover" ] ; then
        if [ -e "vars.sh" ]; then
            error "It looks like this borg was already set up, can only recover from fresh start"
        fi
        RECOVER=1

    elif [ "$1" == "--update-paths" ] || [ "$1" == "--update" ] ; then
        if [ ! -e "vars.sh" ]; then
            error "Can't update, not set up yet"
        fi
        UPDATE=1

    elif [ -n "$1" ] ; then
        error "Unknown arg $1"

    elif [ -e "vars.sh" ]; then
        warn "Error: BORG_DIR $BORG_DIR already looks set up."
        warn "Use \"git clean\" to return it to original state if desired".
        warn "Or specify --update to refresh things from latest git."
        error "Giving up"
    fi
}

# Make a temp dir to work in
TMP=$(mktemp -d)

# Install some cleanup handlers
cleanup()
{
    set +o errexit
    set +o errtrace
    trap - ERR
    ssh -o ControlPath="$TMP"/ssh-control -O exit x >/dev/null 2>&1
    rm -rf -- "$TMP"
}
cleanup_int()
{
    echo
    cleanup
    exit 1
}
trap cleanup 0
trap cleanup_int 1 2 15

msg()
{
    color="$1"
    shift
    echo -ne "\033[1;${color}m===\033[0;${color}m" "$@"
    echo -e "\033[0m"
}
log(){ msg 33 "$@" ; }
notice() { msg 32 "$@" ; }
warn() { msg 31 "$@" ; }
error() { msg 31 "Error:" "$@" ; exit 1 ; }

print_random_key()
{
    dd if=/dev/urandom | tr -dc 'a-zA-Z0-9' | head -c 16
}

generate_keys()
{
    if [ $RECOVER -eq 1 ] ; then
        notice "Recovering configuration in order to use an existing backup"
        read -s -p "Repo key for \"borg ${HOSTNAME}\": " PASS_REPOKEY
        echo
        read -s -p "Again: " PASS_REPOKEY2
        echo
        if [ -z "$PASS_REPOKEY" ] || [ $PASS_REPOKEY != $PASS_REPOKEY2 ] ; then
            error "Bad repo key"
        fi
    else
        PASS_REPOKEY=$(print_random_key)
    fi
    echo "$PASS_REPOKEY" > passphrase
    chmod 600 passphrase
}

# Run a command on the remote host over an existing SSH tunnel
run_ssh_command()
{
    ssh -o ControlPath="$TMP"/ssh-control use-existing-control-tunnel "$@"
}

# Configure SSH key-based login
configure_ssh()
{
    mkdir "$SSH"

    # Create keys
    log "Creating SSH keys"
    ssh-keygen -N "" -t ecdsa \
               -C "backup-appendonly@$HOSTID" -f "$SSH/id_ecdsa_appendonly"
    ssh-keygen -N "" -t ecdsa \
               -C "backup-notify@$HOSTID" -f "$SSH/id_ecdsa_notify"

    # Create config snippets
    log "Creating SSH config and wrapper script"
    cat >> "$SSH/config" <<EOF
User $BACKUP_USER
ControlPath none
ServerAliveInterval 120
Compression no
UserKnownHostsFile $SSH/known_hosts
ForwardX11 no
ForwardAgent no
BatchMode yes
IdentitiesOnly yes
EOF

    # Connect to backup host, using persistent control socket
    log "Connecting to server"
    log "Please enter password; look in Bitwarden for: ssh ${BACKUP_HOST} / ${BACKUP_USER}"
    ssh -F "$SSH/config" -o BatchMode=no -o PubkeyAuthentication=no \
        -o ControlMaster=yes -o ControlPath="$TMP/ssh-control" \
        -o StrictHostKeyChecking=accept-new \
        -p "${BACKUP_PORT}" \
        -f "${BACKUP_USER}@${BACKUP_HOST}" sleep 600
    if ! run_ssh_command true >/dev/null 2>&1 </dev/null ; then
        error "SSH failed"
    fi
    log "Connected to ${BACKUP_USER}@${BACKUP_HOST}:${BACKUP_PORT}"

    # Since we now have an SSH connection, check repo existence
    if [ $RECOVER -eq 0 ] && run_ssh_command "test -e $BACKUP_REPO"; then
        error "$BACKUP_REPO already exists on the server, bailing out"
    elif [ $RECOVER -ne 0 ] && ! run_ssh_command "test -e $BACKUP_REPO"; then
        error "$BACKUP_REPO does NOT exist on the server, can't recover backup config"
    fi

    # Copy SSH keys to the server's authorized_keys file, removing any
    # existing keys with this HOSTID.
    log "Setting up SSH keys on remote host"
    REMOTE_BORG="borg/borg"
    cmd="$REMOTE_BORG serve --restrict-to-repository ~/$BACKUP_REPO"

    keys=".ssh/authorized_keys"
    backup="${keys}.old-$(date +%Y%m%d-%H%M%S)"
    run_ssh_command "mkdir -p .ssh; chmod 700 .ssh; touch $keys"
    run_ssh_command "mv $keys $backup; sed '/@$HOSTID\$/d' < $backup > $keys"
    run_ssh_command "if cmp -s $backup $keys; then rm $backup ; fi"
    run_ssh_command "cat >> .ssh/authorized_keys" <<EOF
command="$cmd --append-only",restrict $(cat "$SSH/id_ecdsa_appendonly.pub")
command="borg/notify.sh",restrict $(cat "$SSH/id_ecdsa_notify.pub")
EOF

    # Test that everything worked
    log "Testing SSH login with new key"
    if ! ssh -F "$SSH/config" -i "$SSH/id_ecdsa_appendonly" -T \
         -p "${BACKUP_PORT}" "${BACKUP_USER}@${BACKUP_HOST}" "$REMOTE_BORG" \
         --version </dev/null ; then
        error "Logging in with a key failed -- is server set up correctly?"
    fi
    log "Remote connection OK!"
}

# Create the repository on the server
create_repo()
{
    log "Creating repo $BACKUP_REPO"
    # Create repo
    $BORG init --make-parent-dirs --encryption repokey
}

# Export keys as HTML page
export_keys()
{
    log "Exporting keys"
    $BORG key export --paper '' key.txt
    chmod 600 key.txt
    cat >>key.txt <<EOF

Repository: ${BORG_REPO}
Passphrase: ${PASS_REPOKEY}
EOF
}

configure_systemd()
{
    TIMER=${SYSTEMD_UNIT}.timer
    SERVICE=${SYSTEMD_UNIT}.service
    TIMER_UNIT=${BORG_DIR}/${TIMER}
    SERVICE_UNIT=${BORG_DIR}/${SERVICE}

    log "Creating systemd files"

    # Choose a time between 1am and 6am based on this hostname
    HASH=$(echo hash of "$HOSTNAME" | sha1sum)
    HOUR=$((0x${HASH:0:8} % 5 + 1))
    MINUTE=$((0x${HASH:8:8} % 6 * 10))
    TIME=$(printf %02d:%02d:00 $HOUR $MINUTE)

    log "Backup time is $TIME"

    cat > "$TIMER_UNIT" <<EOF
[Unit]
Description=Borg backup to ${BACKUP_HOST}

[Timer]
OnCalendar=*-*-* $TIME
Persistent=true

[Install]
WantedBy=timers.target
EOF

    cat >> "$SERVICE_UNIT" <<EOF
[Unit]
Description=Borg backup to ${BACKUP_HOST}

[Service]
Type=simple
ExecStart=${BORG_DIR}/backup.py
Nice=10
IOSchedulingClass=best-effort
IOSchedulingPriority=6
Restart=on-failure
RestartSec=600
EOF

    if [ $RECOVER -eq 1 ] ; then
        log "Partially setting up systemd"
        ln -sfv "${TIMER_UNIT}" /etc/systemd/system
        ln -sfv "${SERVICE_UNIT}" /etc/systemd/system
        systemctl --no-ask-password daemon-reload
        systemctl --no-ask-password stop ${TIMER}
        systemctl --no-ask-password disable ${TIMER}
	warn "Since we're recovering, systemd automatic backups aren't enabled"
	warn "Do something like this to configure automatic backups:"
	echo "     sudo systemctl enable ${TIMER} &&"
        echo "     sudo systemctl start ${TIMER}"
        warn ""
    else
        log "Setting up systemd"
        if (
            ln -sfv "${TIMER_UNIT}" /etc/systemd/system &&
            ln -sfv "${SERVICE_UNIT}" /etc/systemd/system &&
            systemctl --no-ask-password daemon-reload &&
            systemctl --no-ask-password enable ${TIMER} &&
            systemctl --no-ask-password start ${TIMER}
        ); then
            log "Backup timer installed:"
            systemctl list-timers --no-pager ${TIMER}
        else
            warn ""
            warn "Systemd setup failed"
            warn "Do something like this to configure automatic backups:"
            echo "     sudo ln -sfv \"${TIMER_UNIT}\" /etc/systemd/system &&"
            echo "     sudo ln -sfv \"${SERVICE_UNIT}\" /etc/systemd/system &&"
            echo "     sudo systemctl daemon-reload &&"
            echo "     sudo systemctl enable ${TIMER} &&"
            echo "     sudo systemctl start ${TIMER}"
            warn ""
        fi
    fi
}

git_setup()
{
    log "Committing local changes to git"

    if ! git checkout -b "setup-${HOSTNAME}" ||
       ! git add ${SYSTEMD_UNIT}.service ${SYSTEMD_UNIT}.timer vars.sh ||
       ! git commit -a -m "autocommit after initial setup on ${HOSTNAME}" ; then
        warn "Git setup failed; ignoring"
        return
    fi
}

main() {
    if [ $UPDATE -eq 1 ] ; then
        notice "Non-destructively updating paths, variables, and venv..."
        source vars.sh
        set_default_variables
        install_templated_files
        update_paths
        setup_venv
        create_borg_vars

        notice "Testing SSH"
        ssh -F "$SSH/config" -i "$SSH/id_ecdsa_appendonly" \
	     -o BatchMode=no -o StrictHostKeyChecking=ask \
             -p "${BACKUP_PORT}" "${BACKUP_USER}@${BACKUP_HOST}" info >/dev/null

        notice "Testing borg: if host location changed, say 'y' here"
        ${BORG_DIR}/borg.sh info

        notice "Done -- check 'git diff' and verify changes."
        exit 0
    fi

    set_default_variables
    setup_venv
    create_borg_vars
    generate_keys
    configure_ssh
    [ $RECOVER -eq 0 ] && create_repo
    export_keys
    configure_systemd
    install_templated_files
    update_paths
    git_setup

    echo
    if [ $RECOVER -eq 1 ] ; then
        notice "You should be set up with borg pointing to the existing repo now."
        notice "Use commands like these to look at the backup:"
        notice "    sudo ${BORG_DIR}/borg.sh info"
        notice "    sudo ${BORG_DIR}/borg.sh list"
        notice "You'll want to now restore files like ${BORG_DIR}/config.yaml before enabling systemd timers"
    else
        notice "Add this password to Bitwarden:"
        notice ""
        notice "      Name: borg ${HOSTNAME}"
        notice "  Username: repo key"
        notice "  Password: $PASS_REPOKEY"
        notice ""
        notice "Test the backup file list with"
        notice "    sudo ${BORG_DIR}/backup.py --dry-run"
        notice "and make any necessary adjustments to:"
        notice "    ${BORG_DIR}/config.yaml"
    fi

    echo

    echo "All done"
}

parse_args "$@"
main