#!/bin/bash
#
# Copyright 2011-2018 Nicolas Thauvin. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 
#  1. Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.
#  2. Redistributions in binary form must reproduce the above copyright
#     notice, this list of conditions and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
# 
# THIS SOFTWARE IS PROVIDED BY THE AUTHORS ``AS IS'' AND ANY EXPRESS OR
# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
# OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE AUTHORS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
# THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

version="2.3"

# Default configuration
SYSLOG="no"
ARCHIVE_COMPRESS="yes"
ARCHIVE_UNCOMPRESS_BIN=gunzip
ARCHIVE_COMPRESS_SUFFIX=gz

CONFIG_DIR="/etc/pitrery"
CONFIG="pitrery.conf"

GPG_BIN="/usr/bin/gpg"

# Apply an extra level of shell quoting to each of the arguments passed.
# This is necessary for remote-side arguments of ssh (including commands that
# are executed by the remote shell and remote paths for scp and rsync via ssh)
# since they will strip an extra level of quoting off on the remote side.
# This makes it safe for them to include spaces or other special characters
# which should not be interpreted or cause word-splitting on the remote side.
qw() {
    while (( $# > 1 )); do
	printf "%q " "$1"
	shift
    done
    (( $# == 1 )) && printf "%q" "$1"
}

# Message functions
now() {
    [ "$LOG_TIMESTAMP" = "yes" ] && echo "$(date "+%F %T %Z ")"
}

error() {
    echo "$(now)ERROR: $*" 1>&2
    exit 1
}

warn() {
    echo "$(now)WARNING: $*" 1>&2
}

# Script help
usage() {
    echo "$(basename $0) - Restore a WAL segment"
    echo
    echo "usage: `basename $0` [options] xlogfile destination"
    echo "options:"
    echo "    -C conf                Configuration file"
    echo "    -a [[user@]host:]/dir  Place to get the archive"
    echo "    -X                     Do not uncompress"
    echo "    -c command             Uncompression command"
    echo "    -s suffix              Compressed file suffix (ex: gz)"
    echo "    -S                     Send messages to syslog"
    echo "    -f facility            Syslog facility"
    echo "    -t ident               Syslog ident"
    echo "    -T                     Timestamp log messages"
    echo
    echo "    -V                     Display the version and exit"
    echo "    -?                     Print help"
    echo
    exit $1
}

parse_target_uri() {
    local backup_target=$1
    local archive_target=$2

    if [ -n "$backup_target" ]; then
        # Parse the backuptarget into user, host and path
        backup_user="$(echo $backup_target | grep '@' | cut -d'@' -f1 )"
        backup_host="$(echo $backup_target | grep ':' | sed -re 's/(.*):(.*)/\1/' | cut -d'@' -f2 )"
        backup_dir="$(echo $backup_target | sed -re 's/(.*):(.*)/\2/')"
       
    else
        # Fallback to the values from the configuration file
        [ -n "$BACKUP_USER" ] && backup_user=$BACKUP_USER
        [ -n "$BACKUP_HOST" ] && backup_host=$BACKUP_HOST
        [ -n "$BACKUP_DIR" ] && backup_dir=$BACKUP_DIR
    fi

    # Deduce if backup is local
    if [ -z "$backup_host" ]; then
        backup_local="yes"
    else
        backup_local="no"

        # Wrap IPv6 addresses with brackets
        echo $backup_host | grep -qi '^[0123456789abcdef:]*:[0123456789abcdef:]*$' && backup_host="[${backup_host}]"

        # Add a shortcut for ssh/rsync commands
        backup_ssh_target=${backup_user:+$backup_user@}$backup_host
    fi

    if [ -n "$backup_dir" ]; then
        # Ensure the backup directory is an absolute path
        if [ "$backup_local" = "yes" ]; then
            backup_dir="$(readlink -m -- "$backup_dir")"
        else
            backup_dir="$(ssh -n -- "$backup_ssh_target" "readlink -m -- $(qw "$backup_dir")")"
        fi
    fi

    # Parse archive target the same way
    if [ -n "$archive_target" ]; then
        archive_user="$(echo $archive_target | grep '@' | cut -d'@' -f1 )"
        archive_host="$(echo $archive_target | grep ':' | sed -re 's/(.*):(.*)/\1/' | cut -d'@' -f2 )"
        archive_dir="$(echo $archive_target | sed -re 's/(.*):(.*)/\2/')"
    else
        # Fallback to the values of the configuration file. When the
        # path is not provided in the config file, fallback to backup values
        if [ -n "$ARCHIVE_DIR" ]; then
            [ -n "$ARCHIVE_USER" ] && archive_user=$ARCHIVE_USER
            [ -n "$ARCHIVE_HOST" ] && archive_host=$ARCHIVE_HOST
            archive_dir=$ARCHIVE_DIR
        else
            archive_user="$backup_user"
            archive_host="$backup_host"
            # avoid trying to create directory in /
            [ -n "$backup_dir" ] && archive_dir="$backup_dir/archived_xlog"
        fi
    fi

    # Deduce if archives are local
    if [ -z "$archive_host" ]; then
	archive_local="yes"
    else
	archive_local="no"

        # Wrap IPv6 addresses with brackets
        echo $archive_host | grep -qi '^[0123456789abcdef:]*:[0123456789abcdef:]*$' && archive_host="[${archive_host}]"

        # Add a shortcut for ssh/rsync commands
        archive_ssh_target=${archive_user:+$archive_user@}$archive_host
    fi

    [ -n "$archive_dir" ] || error "missing archive directory"

    # Ensure the achives directory is an absolute path
    if [ "$archive_local" = "yes" ]; then
        archive_dir="$(readlink -m -- "$archive_dir")"
    else
        archive_dir="$(ssh -n -- "$archive_ssh_target" "readlink -m -- $(qw "$archive_dir")")"
    fi

}

# CLI processing
while getopts "C:a:Xc:s:Sf:t:TV?" opt; do
    case $opt in
	C) CONFIG=$OPTARG;;
        a) archive_path="$OPTARG";;
	X) CLI_ARCHIVE_COMPRESS="no";;
	c) CLI_ARCHIVE_UNCOMPRESS_BIN=$OPTARG;;
	s) CLI_ARCHIVE_COMPRESS_SUFFIX=$OPTARG;;
	S) CLI_SYSLOG="yes";;
	f) CLI_SYSLOG_FACILITY=$OPTARG;;
	t) CLI_SYSLOG_IDENT=$OPTARG;;
	T) CLI_LOG_TIMESTAMP="yes";;
        V) echo "restore_xlog (pitrery) $version"; exit 0;;
	"?") usage 1;;
	*) error "Unknown error while processing options";;
    esac
done

# Check if the config option is a path or just a name in the
# configuration directory.  Prepend the configuration directory and
# .conf when needed.
if [[ $CONFIG != */* ]]; then
    CONFIG="$CONFIG_DIR/$(basename -- "$CONFIG" .conf).conf"
fi

# Load configuration file
if [ -f "$CONFIG" ]; then
    . "$CONFIG"
fi

# Override configuration with cli options
[ -n "$CLI_ARCHIVE_COMPRESS" ] && ARCHIVE_COMPRESS=$CLI_ARCHIVE_COMPRESS
[ -n "$CLI_ARCHIVE_UNCOMPRESS_BIN" ] && ARCHIVE_UNCOMPRESS_BIN=$CLI_ARCHIVE_UNCOMPRESS_BIN
[ -n "$CLI_ARCHIVE_COMPRESS_SUFFIX" ] && ARCHIVE_COMPRESS_SUFFIX=$CLI_ARCHIVE_COMPRESS_SUFFIX
[ -n "$CLI_SYSLOG" ] && SYSLOG=$CLI_SYSLOG
[ -n "$CLI_SYSLOG_FACILITY" ] && SYSLOG_FACILITY=$CLI_SYSLOG_FACILITY
[ -n "$CLI_SYSLOG_IDENT" ] && SYSLOG_IDENT=$CLI_SYSLOG_IDENT
[ -n "$CLI_LOG_TIMESTAMP" ] && LOG_TIMESTAMP=$CLI_LOG_TIMESTAMP

# Redirect output to syslog if configured
if [ "$SYSLOG" = "yes" ]; then
    SYSLOG_FACILITY=${SYSLOG_FACILITY:-local0}
    SYSLOG_IDENT=${SYSLOG_IDENT:-postgres}

    exec 1> >(logger -t "$SYSLOG_IDENT" -p "${SYSLOG_FACILITY}.info")
    exec 2> >(logger -t "$SYSLOG_IDENT" -p "${SYSLOG_FACILITY}.err")
fi

parse_target_uri "" "$archive_path"

# Check input: the name of the xlog file (%f) is needed as well has the target path (%p)
# PostgreSQL gives those two when executing restore_command
xlog=${@:$OPTIND:1}
target_path=${@:$(($OPTIND+1)):1}

if [ -z "$xlog" ] || [ -z "$target_path" ]; then
    error "missing xlog filename and/or target path. Please use %f and %p in restore_command"
fi

# Get the file: use cp when the file is on localhost, scp otherwise
if [ "$archive_local" = "yes" ]; then
    if [ -f "$archive_dir/$xlog" ]; then
        xlog_file="$xlog"
        target_file="$target_path"
    elif [ -f "$archive_dir/${xlog}.$ARCHIVE_COMPRESS_SUFFIX" ]; then
        xlog_file="${xlog}.$ARCHIVE_COMPRESS_SUFFIX"
        target_file="${target_path}.$ARCHIVE_COMPRESS_SUFFIX"
    elif [ -f "$archive_dir/${xlog}.gpg" ]; then
        xlog_file="${xlog}.gpg"
        target_file="${target_path}.gpg"
        ARCHIVE_COMPRESS="no"
    else
	error "could not find $archive_dir/$xlog"
    fi

    if ! cp -- "$archive_dir/$xlog_file" "$target_file"; then
	error "could not copy $archive_dir/$xlog_file to $target_file"
    fi
else
    if ssh -- "${archive_ssh_target}" "test -f $(qw "$archive_dir/$xlog")"; then
        xlog_file="$xlog"
        target_file="$target_path"
    elif ssh -- "${archive_ssh_target}" "test -f $(qw "$archive_dir/${xlog}.$ARCHIVE_COMPRESS_SUFFIX")"; then
        xlog_file="${xlog}.$ARCHIVE_COMPRESS_SUFFIX"
        target_file="${target_path}.$ARCHIVE_COMPRESS_SUFFIX"
    elif ssh -- "${archive_ssh_target}" "test -f $(qw "$archive_dir/${xlog}.gpg")"; then
        xlog_file="${xlog}.gpg"
        target_file="${target_path}.gpg"
        ARCHIVE_COMPRESS="no"
    else
	error "could not find $archive_dir/$xlog on ${archive_host}"
    fi

    if ! scp -- "${archive_ssh_target}:$(qw "$archive_dir/$xlog_file")" "$target_file" >/dev/null; then
	error "could not copy ${archive_host}:$archive_dir/$xlog_file to $target_file"
    fi
fi

# Uncompress the file if needed
if [ "$ARCHIVE_COMPRESS" = "yes" ]; then
    if ! $ARCHIVE_UNCOMPRESS_BIN "$target_file"; then
	error "could not uncompress $target_file"
    fi
fi

# Or decrypt it
if [[ "$target_file" =~ \.gpg$ ]]; then
    if ! "$GPG_BIN" "--batch" "--yes" "--decrypt" "--quiet" -o "$target_path" "$target_file"; then
        error "could not decrypt $target_file"
    fi
    rm -- "$target_file" || warn "unable to remove $target_file after decryption"
fi

exit 0
