#!/bin/bash

# Ensure script is run with root privileges
if [ "$EUID" -ne 0 ]; then
  echo "❌ Please run as root (use sudo)."
  exit 1
fi

CONFIG_FILE="/etc/telegram.conf"
ALERT_SCRIPT="/usr/local/bin/telegram-ssh-alert.sh"
PAM_FILE="/etc/pam.d/sshd"

echo "🚀 Starting Telegram SSH Alert Setup..."

# ----------------------------------------------------------------------
# 1. Setup & Validate /etc/telegram.conf
# ----------------------------------------------------------------------
# Helper to read and source config file cleanly
load_config() {
    if [ -f "$CONFIG_FILE" ]; then
        source "$CONFIG_FILE"
    fi
}

load_config

# Check if config exists and has non-empty variables
if [ -z "$TELEGRAM_CHAT_ID" ] || [ -z "$TELEGRAM_BOT_TOKEN" ]; then
    echo "⚙️ Config file missing or contains incomplete variables."
    
    # Prompt for missing values
    if [ -z "$TELEGRAM_CHAT_ID" ]; then
        read -p "Enter Telegram Chat ID: " TELEGRAM_CHAT_ID
    fi

    if [ -z "$TELEGRAM_BOT_TOKEN" ]; then
        read -p "Enter Telegram Bot Token: " TELEGRAM_BOT_TOKEN
    fi

    # Write or update configuration file
    cat <<EOF > "$CONFIG_FILE"
# Telegram Alert Configuration
TELEGRAM_CHAT_ID="$TELEGRAM_CHAT_ID"
TELEGRAM_BOT_TOKEN="$TELEGRAM_BOT_TOKEN"
EOF
    echo "✅ Updated $CONFIG_FILE with valid credentials."
else
    echo "ℹ️ Valid configuration found in $CONFIG_FILE."
fi

# Secure config file permissions
chown root:root "$CONFIG_FILE"
chmod 600 "$CONFIG_FILE"
echo "🔒 Secured permissions for $CONFIG_FILE (600)."

# ----------------------------------------------------------------------
# 2. Setup /usr/local/bin/telegram-ssh-alert.sh
# ----------------------------------------------------------------------
echo "📝 Writing alert script to $ALERT_SCRIPT..."

cat << 'EOF' > "$ALERT_SCRIPT"
#!/bin/bash

# Only notify on session open (successful logins)
if [ "$PAM_TYPE" != "open_session" ]; then
    exit 0
fi

# Path to your configuration file
CONFIG_FILE="/etc/telegram.conf"

# Check if the configuration file exists and read it
if [ -f "$CONFIG_FILE" ]; then
    source "$CONFIG_FILE"
else
    echo "Telegram config file not found at $CONFIG_FILE" >&2
    exit 1
fi

# Verify required variables are set
if [ -z "$TELEGRAM_CHAT_ID" ] || [ -z "$TELEGRAM_BOT_TOKEN" ]; then
    echo "Missing TELEGRAM_CHAT_ID or TELEGRAM_BOT_TOKEN in $CONFIG_FILE" >&2
    exit 1
fi

# Gather connection information
USER="${PAM_USER:-Unknown}"
RHOST="${PAM_RHOST:-Unknown}"
HOSTNAME="$(hostname)"
DATE="$(date '+%Y-%m-%d %H:%M:%S')"

# Build stylish notification message
MESSAGE="🚨 <b>SSH Access Alert</b>
───────────────
👤 <b>User:</b> <code>$USER</code>
🖥️ <b>Host:</b> <code>$HOSTNAME</code>
🌐 <b>IP Address:</b> <code>$RHOST</code>
⏰ <b>Time:</b> $DATE
───────────────"

# Send alert to Telegram API quietly in the background
curl -s -X POST "https://api.telegram.org/bot${TELEGRAM_BOT_TOKEN}/sendMessage" \
    -d "chat_id=${TELEGRAM_CHAT_ID}" \
    -d "parse_mode=HTML" \
    -d "text=${MESSAGE}" > /dev/null 2>&1 &

exit 0
EOF

chmod +x "$ALERT_SCRIPT"
echo "✅ $ALERT_SCRIPT created and set as executable."

# ----------------------------------------------------------------------
# 3. Hook into /etc/pam.d/sshd
# ----------------------------------------------------------------------
PAM_LINE="session optional pam_exec.so $ALERT_SCRIPT"

if grep -qF "$ALERT_SCRIPT" "$PAM_FILE"; then
    echo "ℹ️ PAM rule already present in $PAM_FILE."
else
    echo "🔧 Adding PAM rule to $PAM_FILE..."
    echo "$PAM_LINE" >> "$PAM_FILE"
    echo "✅ PAM configuration updated."
fi

# ----------------------------------------------------------------------
# 4. Run a live test
# ----------------------------------------------------------------------
echo "🧪 Running test notification..."
sudo PAM_TYPE="open_session" PAM_USER="InstallerTest" PAM_RHOST="127.0.0.1" "$ALERT_SCRIPT"

echo "✨ Setup complete! Check your Telegram app for the test notification."