#!/usr/bin/expect -f

# 设置标题和版本
set script_name "Auto SSH Login"
set version "1.0.1"

# 关闭expect的超时
set timeout -1

# 设置窗口自适应
trap {
    set rows [stty rows]
    set cols [stty columns]
    stty rows $rows columns $cols < $spawn_out(slave,name)
} WINCH

# 定义不需要密码的主机列表
set no_password_hosts {
    "150.223.248*"
    "cwai*"
    "123"
    "172.18*"
    "127.0.0.1"
    "dev-aiplatform-gpu-4*"
}

# 显示帮助信息的函数
proc show_help {} {
    global script_name version
    puts "\n$script_name v$version"
    puts "用法: $script_name 用户名 跳板机地址 端口 跳板机密码 \[开发机地址\] \[开发机密码\] \[命令\]"
    puts "\n参数说明:"
    puts "  用户名        - 跳板机用户名"
    puts "  跳板机地址    - 跳板机IP或域名"
    puts "  端口          - 跳板机SSH端口"
    puts "  跳板机密码    - 跳板机登录密码"
    puts "  开发机地址    - \[可选\] 开发机地址 (用户名@主机名)"
    puts "  开发机密码    - \[可选\] 开发机登录密码"
    puts "  命令          - \[可选\] 登录后要执行的命令\n"
    puts "示例:"
    puts "  ./auto_login user jumper.example.com 22 password"
    puts "  ./auto_login user jumper.example.com 22 password dev@server.com devpassword"
    puts "  ./auto_login user jumper.example.com 22 password dev@server.com \"\" \"cd /home/work\"\n"
    exit 0
}

# 参数检查
if {[llength $argv] < 4} {
    puts "错误: 参数不足"
    show_help
}



# 处理帮助参数
if {[lindex $argv 0] == "-h" || [lindex $argv 0] == "--help"} {
    show_help
}

# 设置参数
set USER [lindex $argv 0]
set JUMPER_HOST [lindex $argv 1]
set JUMPER_PORT [lindex $argv 2]
set JUMPER_PASSWORD [lindex $argv 3]
set HOST [lindex $argv 4]
set PASSWORD [lindex $argv 5]
set CMD [lindex $argv 6]

# 日志输出函数
proc log_info {message} {
    puts "\[INFO\] $message"
}

proc log_error {message} {
    puts "\[ERROR\] $message"
}

# SSH到目标机器的函数
proc ssh_to_target {host password cmd} {
    global JUMPER_HOST no_password_hosts

    # 如果 host 为 '0'，跳过 ssh 连接，但如果 cmd 不为空，仍然执行 cmd
    if {"$host" == "0"} {
        log_info "host 为 '0'，跳过目标主机连接"
        if {"$cmd" != ""} {
            send "$cmd\r"
            log_info "执行命令: $cmd"
        }
        return 1
    }

    # 确定SSH参数
    set ssh_options "-o GSSAPIAuthentication=no"

    # 根据目标主机特性选择SSH命令
    if {[string first "172.17" "$host"] != -1} {
        # 对特定IP段使用不同的SSH参数
        set ssh_cmd "ssh $host"
    } else {
        # 检查主机是否在不需要密码的列表中
        set needs_password 1
        foreach pattern $no_password_hosts {
            if {[string match $pattern "$host"] || [string match $pattern "$JUMPER_HOST"]} {
                set needs_password 0
                break
            }
        }

        set ssh_cmd "ssh $ssh_options $host"
        if {!$needs_password} {
            set password ""
        }
    }

    # 发送SSH命令
    send "$ssh_cmd\r"
    log_info "连接到目标主机: $host"

    # 如果提供了密码则处理密码提示
    if {"$password" != ""} {
        expect {
            -re "password:" { 
                send "$password\r"
                log_info "已发送目标主机密码"
            }
            timeout {
                log_error "等待密码提示超时"
                return 0
            }
        }
    }

    # 等待shell提示符并设置环境
    expect {
        -re {[#\$%>❯]} {
            # 设置环境变量
            send "export LANG=en_US.UTF-8; export LC_ALL=en_US.UTF-8; export LC_CTYPE=en_US.UTF-8\r"

            # 等待命令执行完成
            expect -re {[#\$%>❯]}

            # 执行自定义命令
            if {"$cmd" != ""} {
                send "$cmd\r"
                log_info "执行命令: $cmd"
            }
        }

        eof {
            log_error "连接到 $host 意外关闭"
            exit 1
        }

        timeout {
            log_error "等待 $host 提示符超时"
            exit 1
        }
    }

    return 1
}

# 主程序
log_info "正在连接跳板机: $USER@$JUMPER_HOST:$JUMPER_PORT"
spawn ssh $USER@$JUMPER_HOST -p $JUMPER_PORT

expect {
    -re "Are you sure you want to continue connecting.*" {
        send "yes\r"
        log_info "已确认主机指纹"
        exp_continue
    }

    -re "password:" {
        send "$JUMPER_PASSWORD\r"
        log_info "已发送跳板机密码"
        exp_continue
    }

    eof {
        log_error "SSH连接意外关闭"
        exit 1
    }

    timeout {
        log_error "SSH连接超时"
        exit 1
    }

    -re "@.*(\\\$|#|%|>|❯)" {
        if {"$HOST" != ""} {
            ssh_to_target $HOST $PASSWORD $CMD
        } else {
            if {"$CMD" != ""} {
                send "$CMD\r"
            }
        }
    }
}

interact
exit

#------------------------------
# 使用示例
#alias autologin='expect $HOME/.auto_login'
#alias dev1='autologin root@host 22 password'
#alias dev2='autologin root@host 22 password work@host password "cd /home/work && clear"'
#alias dev3='autologin root@host 22 "密码"'
#alias dev4='autologin user jumper 22 pass host "" "cd /home/work && clear"'
#alias dev5='autologin user jumper 22 pass work@host'
# not ssh to remote target
#alias dev6='autologin user jumper 22 pass "0"