#!/bin/sh

SSH_HOST=${SSH_HOST:-"localhost"}
SSH_PORT=${SSH_PORT:-20022}
SSH_USER=${SSH_USER:-"backup"}
SSH_PASSWORD=${SSH_PASSWORD:-"backup"}
HOME="/home/$SSH_USER"
SSH_DIR="$HOME/.ssh"
key_type="ed25519"

BACKUP_COMPRESSION=${BACKUP_COMPRESSION:-"zstd"}
BACKUP_PASSWORD=${BACKUP_PASSWORD:-""}

if ! id -u "$SSH_USER" >/dev/null 2>&1; then
    echo "Creating user $SSH_USER..."
    adduser -D -s /bin/sh -h "/home/$SSH_USER" "$SSH_USER"
    # Ensure the user is properly initialized in shadow database
    passwd -u "$SSH_USER" 2>/dev/null || true

fi

if [ ! -d "$SSH_DIR" ]; then                                    
    # Generate host keys if they don't exist                    
    mkdir -p $SSH_DIR                                     
    echo "Generating $key_type host key..."                
    ssh-keygen -t "$key_type" -f "$SSH_DIR"/"id_$key_type" -N "" -q
    chmod 600 $SSH_DIR/id_$key_type                                
                                                                   
    # Add default ssh password if not set                          
    echo "$SSH_USER:$SSH_PASSWORD" | chpasswd                      
    chown -R $SSH_USER:$SSH_USER "/home/$SSH_USER"                 
                                                                   
    echo "Host '$SSH_HOST'                                         
    HostName '$SSH_HOST'                                           
    Port '$SSH_PORT'                                               
    User '$SSH_USER'                                               
    PreferredAuthentications publickey                             
    IdentityFile '$SSH_DIR'/id_'$key_type'        
    IdentitiesOnly yes
    StrictHostKeyChecking no                        
    ">> $SSH_DIR/config                           
                                                  
fi                                                

cd $SSH_DIR
rm /root/.ssh/config 2>/dev/null || true
rm /root/.ssh/known_hosts 2>/dev/null || true
echo "Host *                                                               
    IdentityFile '$SSH_DIR'/id_'$key_type'        
    IdentitiesOnly yes
    StrictHostKeyChecking no                        
"> /root/.ssh/config

if [ "$SSH_HOST" != "localhost" ]; then
    SSH_HOST="$(echo $SSH_HOST | base64 -d | jq -r '.[]')"
fi
# creating cycle for backup services
for BACKUP in $( echo -n $SSH_HOST) ; do      
    # check backup client availability
    RESPONSE=$(echo "exit" | timeout 5 telnet $BACKUP $SSH_PORT | grep Connected)
	echo "RESPONSE: $RESPONSE"
    if [ "$RESPONSE" == "" ]; then
        echo "Expected backup client $BACKUP on port $SSH_PORT currently not available"
        continue
    fi
    # set backup client IP addresses and ports                                         
    sshpass -p "$SSH_PASSWORD" ssh -o StrictHostKeyChecking=no -o PreferredAuthentications=password -o ConnectTimeout=5 -p $SSH_PORT $SSH_USER@$BACKUP exit 2>/dev/null
    PASSWORD_AUTH_EXIT_CODE=$?
    echo "PASSWORD_AUTH_EXIT_CODE: $PASSWORD_AUTH_EXIT_CODE"

    if [ $PASSWORD_AUTH_EXIT_CODE -eq 0 ]; then
        echo "SSH password authentication enabled"
        sshpass -p "$SSH_PASSWORD" ssh -p $SSH_PORT -i $SSH_DIR/id_$key_type $SSH_USER@$BACKUP "mkdir -p ~/.ssh && cat >> ~/.ssh/authorized_keys" < $SSH_DIR/id_$key_type.pub
        ssh -i $SSH_DIR/id_$key_type -p $SSH_PORT $SSH_USER@$BACKUP sed -i "s/PasswordAuthentication\ yes/PasswordAuthentication\ no/g" .ssh/server/sshd_config
        sleep 2
    fi
    sshpass -p "$SSH_PASSWORD" ssh -o StrictHostKeyChecking=no -o PreferredAuthentications=password -o BatchMode=yes -o ConnectTimeout=5 -p $SSH_PORT $SSH_USER@$BACKUP exit 2>/dev/null
    PASSWORD_AUTH_EXIT_CODE=$?
    echo "PASSWORD_AUTH_EXIT_CODE: $PASSWORD_AUTH_EXIT_CODE"
    # check ssh connection via pub key
    if [ $PASSWORD_AUTH_EXIT_CODE -ne 0 ]; then
        echo "SSH password authentication is disabled on the server."

        # check borg backup state
        export $BACKUP_PASSWORD        # Replace your current borg check section with:
        # check borg backup state
        CHECK_BACKUP_STATE=$(ssh -i $SSH_DIR/id_$key_type -p $SSH_PORT $SSH_USER@$BACKUP "BORG_PASSPHRASE='$BACKUP_PASSWORD' borg info /backup/ 2> /dev/null")
        if [ -z "$CHECK_BACKUP_STATE" ]; then
            echo "Ready to init borg backup"
            # Initialize borg repository with encryption
            if [ -z "$BACKUP_COMPRESSION" ]; then
                    COMPRESSION=$(echo "--compression $BACKUP_COMPRESSION")
            else
                    COMPRESSION=""
            fi

            ssh -i $SSH_DIR/id_$key_type -p $SSH_PORT $SSH_USER@$BACKUP "BORG_PASSPHRASE='$BORG_PASSPHRASE' borg init --encryption=repokey-blake2 $SSH_DIR/backup/"
            echo "Borg repository initialized with encryption"
        else
            # start backup
            echo "Borg backup is already initialized. Starting backup..."
            # Your backup command here with passphrase
            ssh -i $SSH_DIR/id_$key_type -p $SSH_PORT $SSH_USER@$BACKUP "BACKUP_PASSWORD='$BACKUP_PASSWORD' borg create $COMPRESSION  /backup/::{hostname}-{now} $SSH_DIR/backup/"
        fi
    fi
done