#!/usr/bin/env bash

###############################################################################
# Copyright (c) 2024 Zion Networks UG                                         #
#                                                                             #
# Permission is hereby granted, free of charge, to any person obtaining a     #
# copy of this software and associated documentation files (the "Software"),  #
# to deal in the Software without restriction, including without limitation   #
# the rights to use, copy, modify, merge, publish, distribute, sublicense,    #
# and/or sell copies of the Software, and to permit persons to whom the       #
# Software is furnished to do so, subject to the following conditions:        #
#                                                                             #
# The above copyright notice and this permission notice shall be included in  #
# all copies or substantial portions of the Software.                         #
#                                                                             #
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR  #
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,    #
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE #
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER      #
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,             #
# ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER #
# DEALINGS IN THE SOFTWARE.                                                   #
###############################################################################

# Planned Features:
# - Add support for more distributions
# - Add support for automatic setup (non-interactive)

# SETTINGS - Make sure to adjust these settings to your needs
DEBUG=0 # Set to 1 to enable debug mode

BACKUP_REPO="Backups/$(hostname | tr '[:lower:]' '[:upper:]')" # Repository path on the backup server
BACKUP_ENCRYPTION="none" # Encryption method for the backup repository (default: none)
BACKUP_PASSPHRASE="example" # Make sure to change this to a secure passphrase (will be ignored if encryption is set to none)
BACKUP_DIRS="/etc /var/log /home" # Directories to backup, separated by space
BACKUP_LABEL_PREFIX="BACKUP-" # Prefix for the backup label
BACKUP_CRON_SCHEDULE="0 3 * * *" # Cron schedule for backups (default: daily at 3:00 AM)

BORGMATIC_CONFIG_FILE="/etc/borgmatic/config.yaml" # Path to the borgmatic configuration file - DO NOT CHANGE UNLESS YOU KNOW WHAT YOU ARE DOING

SSH_HOST="127.0.0.1" # Hostname or IP address of the backup server
SSH_PORT="22" # SSH port of the backup server (default: 22)
SSH_USER="root" # SSH user on the backup server
SSH_PASSWORD="12345678" # Only required for ssh-copy-id
SSH_KEY_NAME="borgmatic"
SSH_KEY_TYPE="ed25519" # SSH key type (e.g. rsa, dsa, ecdsa, ed25519), default: ed25519
SSH_KEY_FILE="/root/.ssh/${SSH_KEY_NAME}_${SSH_KEY_TYPE}"

# SCRIPT - DO NOT EDIT

# constants
readonly SCRIPT_NAME="Borgmatic Backup Setup Tool"
readonly SCRIPT_AUTHOR="Zion Networks at admin@zion-networks.de"
readonly SCRIPT_SUPPORT="admin@zion-networks.de"
readonly VERSION="1.0.0"

# logging functions
function inf {
    echo -e "\e[97m[$(date +%H:%M:%S)] \e[1mINF\e[0m $1"
}

function wrn {
    echo -e "\e[93m[$(date +%H:%M:%S)] \e[1mWRN\e[0m $1"
}

function err {
    echo -e "\e[91m[$(date +%H:%M:%S)] \e[1mERR\e[0m $1"
}

function dbg {
    if [ $DEBUG -eq 1 ]; then
        echo -e "\e[94m[$(date +%H:%M:%S)] \e[1mDBG\e[0m \e[3m$1\e[0m"
    fi
}

# logging prompt functions
function infp {
    echo -e -n "\e[97m[$(date +%H:%M:%S)] \e[1mINF\e[0m $1 "
    read -p "" -r

    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        return 1
    else
        return 0
    fi
}

function wrnp {
    echo -e -n "\e[93m[$(date +%H:%M:%S)] \e[1mWRN\e[0m $1 "
    read -p "" -r

    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        return 1
    else
        return 0
    fi
}

function errp {
    echo -e -n "\e[91m[$(date +%H:%M:%S)] \e[1mERR\e[0m $1 "
    read -p "" -r

    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        return 1
    else
        return 0
    fi
}

# logging functions with follow up (like "[HH:MM:SS] INF Installig package ... OK" where OK is a follow up)
function inf_follow {
    if [ $DEBUG -eq 1 ]; then
        echo -e "\e[97m[$(date +%H:%M:%S)] \e[1mINF\e[0m $1"
    else
        echo -e -n "\e[97m[$(date +%H:%M:%S)] \e[1mINF\e[0m $1"
    fi

    dbg "Running ${@:4}"

    ${@:4}
    _r=$?

    if [ $_r -eq 0 ]; then
        if [ $DEBUG -eq 1 ]; then
            echo -e "\e[97m[$(date +%H:%M:%S)] \e[1mINF\e[0m $1: $2"
        else
            echo -e " $2"
        fi
    else
        if [ $DEBUG -eq 1 ]; then
            echo -e "\e[97m[$(date +%H:%M:%S)] \e[1mINF\e[0m $1: $3"
        else
            echo -e " $3"
        fi
    fi

    return $_r
}

function wrn_follow {
    if [ $DEBUG -eq 1 ]; then
        echo -e "\e[93m[$(date +%H:%M:%S)] \e[1mWRN\e[0m $1"
    else
        echo -e -n "\e[93m[$(date +%H:%M:%S)] \e[1mWRN\e[0m $1"
    fi

    _r=$(${@:4})

    if [ $_r -eq 0 ]; then
        if [ $DEBUG -eq 1 ]; then
            echo -e "\e[93m[$(date +%H:%M:%S)] \e[1mWRN\e[0m $1: $2"
        else
            echo -e " $2"
        fi
    else
        if [ $DEBUG -eq 1 ]; then
            echo -e "\e[93m[$(date +%H:%M:%S)] \e[1mWRN\e[0m $1: $3"
        else
            echo -e " $3"
        fi
    fi
}

function err_follow {
    if [ $DEBUG -eq 1 ]; then
        echo -e "\e[91m[$(date +%H:%M:%S)] \e[1mERR\e[0m $1"
    else
        echo -e -n "\e[91m[$(date +%H:%M:%S)] \e[1mERR\e[0m $1"
    fi

    _r=$(${@:4})

    if [ $_r -eq 0 ]; then
        if [ $DEBUG -eq 1 ]; then
            echo -e "\e[91m[$(date +%H:%M:%S)] \e[1mERR\e[0m $1: $2"
        else
            echo -e " $2"
        fi
    else
        if [ $DEBUG -eq 1 ]; then
            echo -e "\e[91m[$(date +%H:%M:%S)] \e[1mERR\e[0m $1: $3"
        else
            echo -e " $3"
        fi
    fi
}

# setup functions
function apt_update {
    exec 3>&1
    status=$(apt-get update 2>&1 | while read -r line; do dbg "$line"; done 1>&3)
    exec 3>&-
    return $status
}

function apt_upgrade {
    exec 3>&1
    status=$(apt-get upgrade -y 2>&1 | while read -r line; do dbg "$line"; done 1>&3)
    exec 3>&-

    return $status
}

function apt_install {
    exec 3>&1
    status=$(apt-get install -y $1 2>&1 | while read -r line; do dbg "$line"; done 1>&3)
    exec 3>&-

    return $status
}

function pipx_install {
    exec 3>&1
    status=$(pipx install $1 2>&1 | while read -r line; do dbg "$line"; done 1>&3)
    exec 3>&-

    return $status
}

function pipx_upgrade {
    exec 3>&1
    status=$(pipx upgrade $1 2>&1 | while read -r line; do dbg "$line"; done 1>&3)
    exec 3>&-

    return $status
}

function apt_is_installed {
    status=$(dpkg -l $1 2>&1 | tee >(while read -r line; do dbg "$line"; done))

    dbg "Status: $status"

    # Check if package is installed based on the output of dpkg -l
    if echo "$status" | grep -q "^ii"; then
        return 0
    else
        return 1
    fi
}

function pipx_is_installed {
    if ! check_command pipx; then
        return 1
    fi

    status=$(pipx list 2>&1 | tee >(while read -r line; do dbg "$line"; done) | grep -q $1)

    return $?
}

function check_command {
    exec 3>&1
    command -v $1 >/dev/null 2>&1
    status=$?
    if [ $status -eq 0 ]; then
        echo "$1 exists" | while read -r line; do dbg "$line"; done 1>&3
    else
        echo "$1 does not exist" | while read -r line; do dbg "$line"; done 1>&3
    fi
    exec 3>&-
    return $status
}

function run_command {
    dbg "Running command: sudo $*"

    exec 3>&1
    output=$(eval "sudo $@" 2>&1)
    status=$?
    echo "$output" | while read -r line; do dbg "$line"; done 1>&3
    exec 3>&-
    return $status
}

function file_exists {
    if [ -f $1 ]; then
        return 0
    else
        return 1
    fi
}

function dir_exists {
    if [ -d $1 ]; then
        return 0
    else
        return 1
    fi
}

function add_crontab {
    line="$BACKUP_CRON_SCHEDULE root PATH=$PATH:/usr/bin:/usr/local/bin /root/.local/bin/borgmatic --verbosity -1 --syslog-verbosity 1 # borgmatic backup job"

    (crontab -l; echo "$line") | crontab -

    return $?
}

function check_crontab {
    crontab -l | grep -q "# borgmatic backup job"

    return $?
}

# main script

echo -e "\e[97m╭─────────────────────────────────────────────────────────────────────────────╮"
echo -e "│                                                                             │"
echo -e "│                         \e[1m$SCRIPT_NAME\e[0m                         │"
echo -e "│                                   \e[1mv$VERSION\e[0m                                    │"
echo -e "│                                                                             │"
echo -e "╰─────────────────────────────────────────────────────────────────────────────╯\e[0m"

inf "For support please contact $SCRIPT_AUTHOR"

if [ $DEBUG -eq 1 ]; then
    dbg "Debug mode is enabled. This will print additional information."
fi

# check if script is running as root
if [ "$EUID" -ne 0 ]; then
    err "Please run this script as root"
    exit 1
fi

# check if script is running on a supported OS (Ubuntu or Debian)
if [ ! -f /etc/os-release ]; then
    err "This script only supports Ubuntu and Debian"
    exit 1
fi

wrn "This script will install \e[1mborgbackup\e[0m, \e[1mpipx\e[0m"
wrn "and \e[1msshpass\e[0m from the official repositories as well"
wrn "as \e[1mborgmatic\e[0m from the official PyPI repository."

if ! wrnp "Do you want to continue? [y/N]"; then
    inf "Aborted."
    exit 1
fi

if ! check_command apt-get; then
    err "apt-get is not installed. This script only supports Ubuntu and Debian."
    exit 1
fi

if ! check_command dpkg; then
    err "dpkg is not installed. This script only supports Ubuntu and Debian."
    exit 1
fi

if ! inf_follow "Updating package lists..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" apt_update; then
    exit 1
fi

if ! check_command sshpass; then
    if ! inf_follow "Installing sshpass..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" apt_install sshpass; then
        err "Installation of sshpass failed!"
        exit 1
    fi
else
    if ! inf_follow "Upgrading sshpass..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" apt_upgrade sshpass; then
        err "Upgrade of sshpass failed!"
        exit 1
    fi
fi

if infp "Do you also want to upgrade all installed packages? [y/N]"; then
    if ! inf_follow "Upgrading packages..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" apt_upgrade; then
        exit 1
    fi
fi

if ! inf_follow "Checking if borgbackup is installed..." "\e[1;32mYes\e[0m" "\e[1;31mNo\e[0m" apt_is_installed borgbackup; then
    if ! inf_follow "Installing borgbackup..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" apt_install borgbackup; then
        err "Installation of borgbackup failed!"
        exit 1
    fi

    if ! inf_follow "Validating ..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" check_command borg; then
        err "Installation of borgbackup failed!"
        exit 1
    fi
fi

if ! inf_follow "Checking if pipx is installed..." "\e[1;32mYes\e[0m" "\e[1;31mNo\e[0m" apt_is_installed pipx; then
    if ! inf_follow "Installing pipx..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" apt_install pipx; then
        err "Installation of pipx failed!"
        exit 1
    fi
fi

if ! inf_follow "Checking if borgmatic is installed..." "\e[1;32mYes\e[0m" "\e[1;31mNo\e[0m" pipx_is_installed borgmatic; then
    if ! inf_follow "Installing borgmatic..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" pipx_install borgmatic; then
        err "Installation of borgmatic failed!"
        exit 1
    fi
else
    if ! inf_follow "Upgrading borgmatic..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" pipx_upgrade borgmatic; then
        err "Upgrade of borgmatic failed!"
        exit 1
    fi
fi

if ! inf_follow "Checking for ~/.ssh directory..." "\e[1;32mYes\e[0m" "\e[1;31mNo\e[0m" dir_exists ~/.ssh; then
    if ! inf_follow "Creating ~/.ssh directory..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" mkdir -p ~/.ssh; then
        err "Failed to create ~/.ssh directory."
        err "Please make sure you have the required permissions."
        exit 1
    fi
fi

if ! inf_follow "Checking for default ssh key at $SSH_KEY_FILE..." "\e[1;32mYes\e[0m" "\e[1;31mNo\e[0m" file_exists $SSH_KEY_FILE; then
    inf "Generating a new ssh key pair..."

    ssh-keygen -t $SSH_KEY_TYPE -C "borgmatic" -f $SSH_KEY_FILE -N "" -q > /dev/null 2>&1

    if [ $? -eq 0 ]; then
        inf "Successfully generated a new ssh key pair."
    else
        err "Failed to generate a new ssh key pair."
        err "If the error persists, please contact the support at $SCRIPT_SUPPORT."
        exit 1
    fi
fi

if ! inf_follow "Copying ssh key to $SSH_HOST..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" run_command sshpass -p "$SSH_PASSWORD" ssh-copy-id -i "$SSH_KEY_FILE" -p "$SSH_PORT" -o StrictHostKeyChecking=no "$SSH_USER@$SSH_HOST"; then
    err "Failed to copy ssh key to $SSH_HOST."
    err "Please validate your ssh password and and host settings and try again."
    err "If the error persists, please contact the support at $SCRIPT_SUPPORT."
    exit 1
else
    if ! inf_follow "Checking if ssh key is working..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" run_command ssh -i "$SSH_KEY_FILE" -p "$SSH_PORT" "$SSH_USER@$SSH_HOST" exit; then
        err "Failed to connect to $SSH_HOST using the ssh key."
        err "Please validate you have the correct ssh key and host settings."
        err "If the error persists, please contact the support at $SCRIPT_SUPPORT."
        exit 1
    fi
fi

if ! inf_follow "Checking if remote end has borg installed..." "\e[1;32mYes\e[0m" "\e[1;31mNo\e[0m" run_command ssh -i "$SSH_KEY_FILE" -p "$SSH_PORT" "$SSH_USER@$SSH_HOST" borg --version; then
    if wrnp "The remote end does not have borg installed. Do you want to install it now? [y/N]"; then
        if ! inf_follow "Updating package repositories on remote end" "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" run_command ssh -i "$SSH_KEY_FILE" -p "$SSH_PORT" "$SSH_USER@$SSH_HOST" apt-get update; then
            err "Failed to update package repositories on remote end."
            err "Please make sure you have the required access rights."
            exit 1
        fi
        
        if ! inf_follow "Installing borg on remote end..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" run_command ssh -i "$SSH_KEY_FILE" -p "$SSH_PORT" "$SSH_USER@$SSH_HOST" apt-get install -y borgbackup; then
            err "Failed to install borg on remote end."
            err "Please make sure you have the required access rights"
            exit 1
        fi
    else
        err "Please make sure borg is installed on the remote end."
        err "This is required in order to do backups using borgmatic."
        exit 1
    fi
fi

inf "Setting up borgmatic configuration..."

# check if directory path exists
if ! inf_follow "Checking for borgmatic configuration directory..." "\e[1;32mYes\e[0m" "\e[1;31mNo\e[0m" dir_exists $(dirname $BORGMATIC_CONFIG_FILE); then
    if ! inf_follow "Creating borgmatic configuration directory..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" mkdir -p $(dirname $BORGMATIC_CONFIG_FILE); then
        err "Failed to create borgmatic configuration directory."
        err "If the error persists, please contact the support at $SCRIPT_SUPPORT."
        exit 1
    fi
fi

cat > $BORGMATIC_CONFIG_FILE <<EOF

ssh_command: ssh -i $SSH_KEY_FILE -p $SSH_PORT

source_directories:
$(echo $BACKUP_DIRS | tr ' ' '\n' | sed -e 's/^/- /')

repositories:
  - path: ssh://$SSH_USER@$SSH_HOST:$SSH_PORT/$BACKUP_REPO
    label: $BACKUP_LABEL_PREFIX$(hostname | tr '[:lower:]' '[:upper:]')

keep_daily: 7
keep_weekly: 4
keep_monthly: 12
keep_yearly: 1

checks:
  - name: repository
  - name: archives
    frequency: 4 weeks

EOF

if [ $? -eq 0 ]; then
    inf "Successfully created borgmatic configuration file."
else
    err "Failed to create borgmatic configuration file."
    err "If the error persists, please contact the support at $SCRIPT_SUPPORT."
    exit 1
fi

if ! inf_follow "Validating borgmatic configuration..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" run_command /root/.local/pipx/venvs/borgmatic/bin/borgmatic config validate; then
    err "Validation of borgmatic configuration failed."
    err "Please check the configuration file at $BORGMATIC_CONFIG_FILE."
    err "If the error persists, please contact the support at $SCRIPT_SUPPORT."
    exit 1
fi

if [ "$BACKUP_ENCRYPTION" != "none" ]; then
    if ! inf_follow "Setting up borgmatic repository without encryption..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" run_command /root/.local/pipx/venvs/borgmatic/bin/borgmatic init --make-parent-dirs --encryption=none; then
        err "Failed to set up borgmatic repository without encryption."
        err "If the error persists, please contact the support at $SCRIPT_SUPPORT."
        exit 1
    fi
else
    if ! inf_follow "Setting up borgmatic repository with encryption..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" run_command BORG_PASSPHRASE=$BACKUP_PASSPHRASE /root/.local/pipx/venvs/borgmatic/bin/borgmatic init --make-parent-dirs --encryption=$BACKUP_ENCRYPTION; then
        err "Failed to set up borgmatic repository with encryption."
        err "If the error persists, please contact the support at $SCRIPT_SUPPORT."
        exit 1
    fi
fi

if ! inf_follow "Checking for cronjob..." "\e[1;32mYes\e[0m" "\e[1;31mNo\e[0m" check_crontab; then
    if ! inf_follow "Adding borgmatic cron job..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" add_crontab; then
        err "Failed to add borgmatic cron job."
        err "If the error persists, please contact the support at $SCRIPT_SUPPORT."
        exit 1
    else
        if ! inf_follow "Validating borgmatic cron job..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" check_crontab; then
            err "Validation of borgmatic cron job failed."
            err "Please check the cron job with crontab -l."
            err "If the error persists, please contact the support at $SCRIPT_SUPPORT."
            exit 1
        fi
    fi
fi

if ! infp "Do you want to run the first backup now? [y/N]"; then
    inf "Setup completed."
    exit 0
fi

if ! inf_follow "Running first backup..." "\e[1;32mOK\e[0m" "\e[1;31mFAILED\e[0m" run_command /root/.local/pipx/venvs/borgmatic/bin/borgmatic create --verbosity 1 --list --stats; then
    err "Failed to run first backup."
    err "If the error persists, please contact the support at $SCRIPT_SUPPORT."
    exit 1

else
    inf "Your first backup has been created successfully!"
fi

inf "Setup completed."

exit 0