#!/bin/bash
# SSH稳定性优化脚本
# 用于优化SSH服务配置，保证连接稳定性

set -e

SSHD_CONFIG="/etc/ssh/sshd_config"
BACKUP_FILE="/etc/ssh/sshd_config.bak.$(date +%Y%m%d_%H%M%S)"

echo "=========================================="
echo "SSH Stability Optimization Script"
echo "=========================================="

# 备份原配置
if [ -f "$SSHD_CONFIG" ]; then
    echo "[1/5] Backing up SSH config to: $BACKUP_FILE"
    cp "$SSHD_CONFIG" "$BACKUP_FILE"
else
    echo "[ERROR] SSH config file not found: $SSHD_CONFIG"
    exit 1
fi

# 优化SSH配置
echo "[2/5] Optimizing SSH configuration..."

# 检查并添加配置（避免重复）
optimize_ssh_config() {
    local config_key="$1"
    local config_value="$2"
    
    # 检查配置是否已存在（包括注释掉的）
    if grep -qE "^#?${config_key}\s" "$SSHD_CONFIG"; then
        # 存在则修改（包括注释的配置）
        sed -i "s|^#\?${config_key}\s.*|${config_key} ${config_value}|g" "$SSHD_CONFIG"
        echo "  - Updated: ${config_key} ${config_value}"
    else
        # 不存在则添加
        echo "${config_key} ${config_value}" >> "$SSHD_CONFIG"
        echo "  - Added: ${config_key} ${config_value}"
    fi
}

# 核心稳定性优化
optimize_ssh_config "PermitRootLogin" "yes"
optimize_ssh_config "ClientAliveInterval" "300"
optimize_ssh_config "ClientAliveCountMax" "3"
optimize_ssh_config "TCPKeepAlive" "yes"
optimize_ssh_config "LoginGraceTime" "60"
optimize_ssh_config "MaxStartups" "10:30:100"
optimize_ssh_config "UseDNS" "no"
optimize_ssh_config "GSSAPIAuthentication" "no"
optimize_ssh_config "PermitUserEnvironment" "yes"

# 性能优化
optimize_ssh_config "MaxSessions" "10"
optimize_ssh_config "MaxAuthTries" "6"
optimize_ssh_config "PubkeyAuthentication" "yes"
optimize_ssh_config "PasswordAuthentication" "yes"

# 安全优化（不影响稳定性）
optimize_ssh_config "Protocol" "2"
optimize_ssh_config "LogLevel" "INFO"
optimize_ssh_config "PermitEmptyPasswords" "no"
optimize_ssh_config "X11Forwarding" "no"
optimize_ssh_config "AllowAgentForwarding" "yes"
optimize_ssh_config "AllowTcpForwarding" "yes"

echo "[3/5] Testing SSH configuration..."
if sshd -t 2>/dev/null; then
    echo "  - Configuration test passed"
else
    echo "  - [WARNING] Configuration test failed, restoring backup"
    cp "$BACKUP_FILE" "$SSHD_CONFIG"
    exit 1
fi

# 重启SSH服务
echo "[4/5] Restarting SSH service..."
if command -v systemctl &> /dev/null && systemctl is-active sshd &> /dev/null; then
    systemctl restart sshd
    sleep 2
    if systemctl is-active sshd &> /dev/null; then
        echo "  - SSH service restarted successfully (systemctl)"
    else
        echo "  - [ERROR] SSH service failed to restart"
        exit 1
    fi
elif command -v service &> /dev/null; then
    service ssh restart
    sleep 2
    if pgrep -x "sshd" > /dev/null; then
        echo "  - SSH service restarted successfully (service)"
    else
        echo "  - [ERROR] SSH service failed to restart"
        exit 1
    fi
else
    echo "  - [WARNING] Could not restart SSH service automatically"
fi

# 验证优化结果
echo "[5/5] Verifying optimization..."
echo ""
echo "Optimized SSH Configuration:"
echo "=========================================="
grep -E "^[^#]*(ClientAliveInterval|ClientAliveCountMax|TCPKeepAlive|LoginGraceTime|MaxStartups|UseDNS)" "$SSHD_CONFIG" || true

echo ""
echo "=========================================="
echo "SSH optimization completed successfully!"
echo "Backup file: $BACKUP_FILE"
echo "=========================================="
echo ""
echo "Recommended next steps:"
echo "  1. Test SSH connection: ssh -v user@localhost"
echo "  2. Monitor logs: tail -f /var/log/auth.log"
echo "  3. Verify keepalive: netstat -an | grep :22"
echo ""