#!/bin/bash

# ssh-notify.sh
# part of ssh-notify from server-toolset
# 2025 © Yigid BALABAN <fyb@fybx.dev>

# --- Configuration ---
CONFIG_FILE="/etc/ssh-notify/config.conf"
HOSTNAME=$(hostname -f)

# --- Script Variables ---
LOGIN_USER=""
REMOTE_HOST=""
MODE="PAM"
TEST_TYPE="both"

# --- Helper Functions ---
usage() {
    echo
    echo "Usage:"
    echo "  (As PAM module): Called automatically by PAM on SSH session start"
    echo "  (For Testing):    sudo $0 -t[=email|telegram] <test_user> <test_remote_host>"
    echo "      -t, --test             : Test both email and Telegram"
    echo "      -t email, --test=email : Test email only"
    echo "      -t telegram, --test=telegram : Test Telegram only"
    echo "      Example: sudo $0 -t testuser 192.168.1.100"
    echo "      Example: sudo $0 -t email testuser 192.168.1.100"
    echo "      Example: sudo $0 --test=telegram testuser 192.168.1.100"
    exit 1
}

log_message() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] [$1] $2" >> "$LOG_FILE"
}

# --- Argument Parsing & Mode Detection ---
case "$1" in
    -t|--test)
        MODE="TEST"
        TEST_TYPE="both"
        shift
        if [[ "$1" == "email" ]]; then
            TEST_TYPE="email"
            shift
        elif [[ "$1" == "telegram" ]]; then
            TEST_TYPE="telegram"
            shift
        fi
        ;;
    -t=email|--test=email)
        MODE="TEST"
        TEST_TYPE="email"
        shift
        ;;
    -t=telegram|--test=telegram)
        MODE="TEST"
        TEST_TYPE="telegram"
        shift
        ;;
    *)
        MODE="PAM"
        ;;
esac

if [[ "$MODE" == "TEST" ]]; then
    if [[ $# -lt 2 ]]; then
        echo "Error: Test mode requires <test_user> and <test_remote_host> arguments."
        usage
    fi
    LOGIN_USER="$1"
    REMOTE_HOST="$2"
    echo "--- Running in TEST mode ($TEST_TYPE) ---"
    echo "User: $LOGIN_USER"
    echo "Host: $REMOTE_HOST"
else
    # --- PAM Mode Logic ---
    if [[ "$PAM_TYPE" != "open_session" || -z "$PAM_USER" || -z "$PAM_RHOST" ]]; then
        exit 0
    fi
    LOGIN_USER="$PAM_USER"
    REMOTE_HOST="$PAM_RHOST"
fi

# --- Load Configuration (needed for both modes) ---
if [[ -f "$CONFIG_FILE" ]]; then
    source "$CONFIG_FILE"
else
    ERR_MSG="ssh-notify Error: Configuration file $CONFIG_FILE not found."
    echo "$ERR_MSG" | systemd-cat -p err -t 'ssh-notify'
    echo "$ERR_MSG" >&2
    exit 1
fi

# Ensure required config variables are set
if [[ -z "$EMAIL_RECIPIENT" || -z "$EMAIL_API_ENDPOINT" || -z "$TELEGRAM_BOT_TOKEN" || -z "$TELEGRAM_CHAT_ID" || -z "$LOG_FILE" ]]; then
    ERR_MSG="ssh-notify Error: One or more required variables are missing in $CONFIG_FILE."
    echo "$ERR_MSG" | systemd-cat -p err -t 'ssh-notify'
    echo "$ERR_MSG" >&2
    exit 1
fi

# Prepare log directory and file
LOG_DIR=$(dirname "$LOG_FILE")
mkdir -p "$LOG_DIR"
if [[ ! -f "$LOG_FILE" ]]; then
     touch "$LOG_FILE"
     # Set permissions only if we create it
     chown root:root "$LOG_FILE" # Or another appropriate user/group if needed
     chmod 640 "$LOG_FILE"     # Restrict write access
fi

# Log the mode of operation
log_message INFO "Script triggered. Mode: $MODE. User: '$LOGIN_USER'. Host: '$REMOTE_HOST'."

# --- Prepare Notification Details ---
TIMESTAMP=$(date '+%Y-%m-%d %H:%M:%S %Z')
SUBJECT="[${MODE}] SSH Login on $HOSTNAME: $LOGIN_USER" # Add mode to subject
MESSAGE_TEXT="[${MODE}] SSH login detected on $HOSTNAME: User '$LOGIN_USER' from '$REMOTE_HOST' at $TIMESTAMP"
MESSAGE_HTML="<b>[${MODE}] SSH Login Alert</b>%0AServer: $HOSTNAME%0AUser: $LOGIN_USER%0AFrom: $REMOTE_HOST%0ATime: $TIMESTAMP"

# --- Notification Functions ---
send_email() {
    log_message INFO "Attempting to send email to $EMAIL_RECIPIENT for user $LOGIN_USER from $REMOTE_HOST"
    JSON_PAYLOAD=$(jq -n \
        --arg subject "$SUBJECT" \
        --arg text "$MESSAGE_TEXT" \
        --arg recipient "$EMAIL_RECIPIENT" \
        '{subject: $subject, text: $text, recipient: $recipient}')

    RESPONSE_CODE=$(curl -s -o /dev/null -w "%{http_code}" \
        -H "Content-Type: application/json" \
        -d "$JSON_PAYLOAD" \
        -L "$EMAIL_API_ENDPOINT")

    if [[ "$RESPONSE_CODE" -eq 200 ]]; then
        log_message INFO "Email alert sent successfully. Response code: $RESPONSE_CODE"
    else
        log_message ERROR "Failure sending email alert. Response code: $RESPONSE_CODE"
    fi
}

send_telegram() {
    log_message INFO "Attempting to send Telegram message to chat $TELEGRAM_CHAT_ID for user $LOGIN_USER from $REMOTE_HOST"
    TELEGRAM_API_URL="https://api.telegram.org/bot${TELEGRAM_BOT_TOKEN}/sendMessage"

    RESPONSE=$(curl -s -X POST "$TELEGRAM_API_URL" \
        -d chat_id="$TELEGRAM_CHAT_ID" \
        -d text="$MESSAGE_HTML" \
        -d parse_mode="HTML")

    if echo "$RESPONSE" | grep -q '"ok":true'; then
        log_message INFO "Telegram alert sent successfully."
    else
        ERROR_DESC=$(echo "$RESPONSE" | jq -r .description)
        log_message ERROR "Failure sending Telegram alert. Response: $ERROR_DESC"
    fi
}

# --- Execute Notifications in Background ---
(
    if [[ "$MODE" == "TEST" ]]; then
        [[ "$TEST_TYPE" == "both" || "$TEST_TYPE" == "email" ]] && send_email
        [[ "$TEST_TYPE" == "both" || "$TEST_TYPE" == "telegram" ]] && send_telegram
    else
        send_email
        send_telegram
    fi
) &

log_message INFO "Notification process forked for user $LOGIN_USER from $REMOTE_HOST"

if [[ "$MODE" == "TEST" ]]; then
    echo "--- Test notifications triggered in background ---"
    echo "Check $LOG_FILE for details."
fi

exit 0