基于Go实现Ubuntu远程Shell捕获及sudo密码按需自动输入

问题概述 #

我想用Go编写自动化脚本，远程连接Ubuntu服务器，实现类似终端的功能，捕获Shell的所有输入、输出和错误信息。目前的实现无法处理sudo命令的密码输入场景：需要仅当终端出现密码请求提示（比如"Enter password:"）时，自动输入硬编码密码，而不是每次通过stdinpipe发送密码。具体场景：

• 执行sudo command，终端提示"Enter password:"时，自动输入密码（当前无法实现该条件判断）
• 执行sudo command，无密码请求时则无需发送密码

原实现代码 #

package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"strings"

	"golang.org/x/crypto/ssh"
)

func main() {
	fmt.Println("Started")
	// Define SSH connection parameters
	host := "host_name"
	port := "port"
	user := "user"
	password := "password"

	// Establish SSH connection
	client, err := connectToSSH(host, port, user, password)
	if err != nil {
		log.Fatalf("Failed to connect to SSH server: %v", err)
	}
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		log.Fatalf("Failed to get New Session: %v", err)
	}
	defer session.Close()

	// Set up standard input/output/error
	stdin, err := session.StdinPipe()
	if err != nil {
		log.Fatalf("Failed to create stdin pipe: %v", err)
	}

	stdout, err := session.StdoutPipe()
	if err != nil {
		log.Fatalf("Failed to create stdout pipe: %v", err)
	}

	stderr, err := session.StderrPipe()
	if err != nil {
		log.Fatalf("Failed to create stderr pipe: %v", err)
	}

	defer stdin.Close()

	outSC := bufio.NewScanner(stdout)
	errSC := bufio.NewScanner(stderr)

	// Start goroutine to handle output
	err = session.Shell()
	if err != nil {
		log.Fatalf("Failed to connect to SSH Shell: %v", err)
	}
	runCommand(stdin, "pwd")
	runCommand(stdin, "sudo su")
	runCommand(stdin, "pwd")

	for outSC.Scan() || errSC.Scan() {
		fmt.Println(outSC.Text())
		fmt.Println(errSC.Text())
		if strings.Contains(outSC.Text(), "sudo") && strings.Contains(outSC.Text(), "password") {
			runCommand(stdin, "password")
		}
	}

	// Wait for the session to finish
	err = session.Wait()
	if err != nil {
		log.Fatalf("Session returned error: %v", err)
	}
}

func connectToSSH(host, port, user, password string) (*ssh.Client, error) {
	// Create SSH client config with password
	config := &ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{
			ssh.Password(password),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), // Insecure; use known_hosts in production
		BannerCallback:  ssh.BannerDisplayStderr(),
	}

	// Connect to SSH server
	client, err := ssh.Dial("tcp", host+":"+port, config)
	if err != nil {
		return nil, err
	}
	return client, nil
}

func runCommand(stdin io.WriteCloser, cmd string) error {
	fmt.Println()
	fmt.Println(cmd)
	_, err := stdin.Write([]byte(cmd + "\n")) // Add a newline character
	if err != nil {
		log.Println("Failed to write to stdin:", err)
		return err
	}
	return nil
}

解决方案 #

原代码的核心问题是串行监听输出的逻辑会阻塞，导致无法及时捕获密码提示。下面是修改后的实现，通过并行监听输出流来实时响应密码请求：

修改后的代码 #

package main

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"strings"

	"golang.org/x/crypto/ssh"
)

func main() {
	fmt.Println("Started")
	// 配置参数
	host := "host_name"
	port := "port"
	user := "user"
	sshPassword := "password"
	sudoPassword := sshPassword // 若sudo密码与SSH密码不同，可单独设置

	// 建立SSH连接
	client, err := connectToSSH(host, port, user, sshPassword)
	if err != nil {
		log.Fatalf("SSH连接失败: %v", err)
	}
	defer client.Close()

	session, err := client.NewSession()
	if err != nil {
		log.Fatalf("创建Session失败: %v", err)
	}
	defer session.Close()

	// 建立输入输出管道
	stdin, err := session.StdinPipe()
	if err != nil {
		log.Fatalf("创建stdin管道失败: %v", err)
	}
	defer stdin.Close()

	stdout, err := session.StdoutPipe()
	if err != nil {
		log.Fatalf("创建stdout管道失败: %v", err)
	}

	stderr, err := session.StderrPipe()
	if err != nil {
		log.Fatalf("创建stderr管道失败: %v", err)
	}

	// 并行监听stdout，实时检查密码提示
	go func() {
		scanner := bufio.NewScanner(stdout)
		for scanner.Scan() {
			line := scanner.Text()
			fmt.Println("[STDOUT]", line)
			// 匹配常见的sudo密码提示
			if strings.Contains(line, "[sudo] password for") || strings.Contains(line, "Enter password:") {
				_, err := stdin.Write([]byte(sudoPassword + "\n"))
				if err != nil {
					log.Printf("发送sudo密码失败: %v", err)
				}
			}
		}
		if err := scanner.Err(); err != nil {
			log.Printf("读取stdout出错: %v", err)
		}
	}()

	// 并行监听stderr，部分系统的sudo提示会输出到stderr
	go func() {
		scanner := bufio.NewScanner(stderr)
		for scanner.Scan() {
			line := scanner.Text()
			fmt.Println("[STDERR]", line)
			if strings.Contains(line, "[sudo] password for") || strings.Contains(line, "Enter password:") {
				_, err := stdin.Write([]byte(sudoPassword + "\n"))
				if err != nil {
					log.Printf("发送sudo密码失败: %v", err)
				}
			}
		}
		if err := scanner.Err(); err != nil {
			log.Printf("读取stderr出错: %v", err)
		}
	}()

	// 启动Shell会话
	err = session.Shell()
	if err != nil {
		log.Fatalf("启动Shell失败: %v", err)
	}

	// 发送测试命令
	runCommand(stdin, "pwd")
	runCommand(stdin, "sudo su")
	runCommand(stdin, "pwd")

	// 等待会话结束
	err = session.Wait()
	if err != nil {
		log.Fatalf("会话执行出错: %v", err)
	}
}

func connectToSSH(host, port, user, password string) (*ssh.Client, error) {
	config := &ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{
			ssh.Password(password),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), // 生产环境请替换为已知主机密钥验证逻辑
		BannerCallback:  ssh.BannerDisplayStderr(),
	}

	client, err := ssh.Dial("tcp", host+":"+port, config)
	if err != nil {
		return nil, err
	}
	return client, nil
}

func runCommand(stdin io.WriteCloser, cmd string) error {
	fmt.Println("\n[INPUT]", cmd)
	_, err := stdin.Write([]byte(cmd + "\n"))
	if err != nil {
		log.Println("写入命令失败:", err)
		return err
	}
	return nil
}

关键修改点 #

1. 并行监听输出流：用两个goroutine分别监听stdout和stderr，避免串行逻辑导致的阻塞，确保不会错过密码提示
2. 完善提示匹配规则：同时覆盖[sudo] password for（Ubuntu默认提示）和Enter password:两种常见的sudo密码提示格式
3. 增加输入输出标识：打印时加入[STDOUT]/[STDERR]/[INPUT]前缀，方便区分不同类型的信息
4. 独立配置sudo密码：将sudo密码单独提取，支持SSH登录密码与sudo密码不同的场景

内容的提问来源于stack exchange，提问作者Abhishek Sharma