通过websocket与ssh来实现远程连接虚拟机

糖孑

已于 2024-04-17 14:14:17 修改

阅读量408

点赞数 6

文章标签： go

于 2024-04-17 14:13:36 首次发布

本文链接：https://blog.csdn.net/m0_52087719/article/details/137870901

版权

1.需求分析

基于Go语言实现一个简单的远程连接程序。前端通过websocket网络连接向后端发送请求，后端接收到前端的数据与用户想要连接的虚拟机建立连接。前端<----->后端<------>远程服务器。

2.实现

时间原因就不设计前端代码了，这里我采用的是ApiFox向后端发出websocket请求。

2.1 websocket

websocket是一种持久化连接的协议，这里采用websocket协议作为前后端请求的原因就是它所花费的开销更小。

2.2 代码实现

这里采用Go语言设计（单纯是因为我最近的业务需要，没有其他理由），采用gorilla的websocket库、golang自带的ssh库。

全部代码已开源到Github。

注意下面代码无特殊说明，就是写在main文件下。

1.首先搭建路由框架，监听8080端口即可。

func main() {
	http.HandleFunc("/api/terminal/sshConnect", sshConnect)
	log.Println("WebSocket server started on :8080")
	log.Fatal(http.ListenAndServe(":8080", nil))
}

2.编写实现函数sshConnetc

func sshConnect(w http.ResponseWriter, r *http.Request) {}

2.1首先获取请求中的参数，我们需要获取传递来的虚拟机IP、密码、登录的用户名

host := r.URL.Query().Get("host")
port := r.URL.Query().Get("port")
password := r.URL.Query().Get("password")
username := r.URL.Query().Get("username")

//将数据保存到一个结构体里
vmInfo := model.VMConnectRequest{
    Host:     host,
    Port:     port,
    Username: username,
    Password: password,
}

2.2 配置ssh客户端

在ssh中分为Client与session对象，session会话对象是真正进行远程连接的对象，而Client对象是用来创建session对象的。

HostKeyCallback属性是用来声明如何对密钥进行处理的，如果需要在连接时判断前端传递的密钥是否正确，那我们就可以这个地方编写相关函数来进行判断，这里我们之间忽略密钥验证。

config := &ssh.ClientConfig{
    User: vmInfo.Username,
    Auth: []ssh.AuthMethod{
       ssh.Password(vmInfo.Password),
    },
    HostKeyCallback: ssh.InsecureIgnoreHostKey(),
}

2.3 测试与远程终端连接是否成功

//测试连接
flag, connctID := utils.VerifyConnect(&vmInfo)
if flag != true {
    //todo 返回错误码
    log.Fatalln("测试连接失败")
    return
}

//utils 包下
func VerifyConnect(vmInfo *model.VMConnectRequest) (bool, string) {
	sshClient, _ := CreateSSHClient(vmInfo, ssh.Password(vmInfo.Password))
	defer sshClient.Close()
	if tempSession, err := sshClient.NewSession(); err == nil {
		defer tempSession.Close()
		//判断是否能够连接成功
		if err := tempSession.Run("whoami"); err == nil {
			//连接成功
			//生成唯一的SSHConnectID
			// 生成一个新的 xid
			myID := xid.New()
			// 将 xid 转换为字符串
			connectID := myID.String()
			return true, connectID
		} else {
			//连接不成功，返回错误
			return false, ""
		}
	} else {
		tempSession.Close()
		return false, ""
	}

}

2.4 能够成功连接后，创建 Client对象与Session

//创建客户端与会话
conn, err := ssh.Dial("tcp", vmInfo.Host+":"+vmInfo.Port, config)
if err != nil {
    log.Fatalf("002unable to connect: %v", err)
    //todo 返回错误码
    return
}
defer conn.Close()

session, err := conn.NewSession()
if err != nil {
    log.Fatalf("003unable to create session: %v", err)
    //todo 返回错误码
    return
}
defer session.Close()

2.5 获取ssh的数据管道对象

ssh的数据管道对象可以看作水管，我们用水管将数据链路连通

stdinPipe, err := session.StdinPipe()
if err != nil {
    log.Fatal(err)
    //todo 返回错误码
}
defer stdinPipe.Close()

// Prepare pipes for capturing output
stdoutPipe, err := session.StdoutPipe()
if err != nil {
    log.Fatal(err)
    //todo 返回错误码
}
stderrPipe, err := session.StderrPipe()
if err != nil {
    log.Fatal(err)
    //todo 返回错误码
}

2.6 将本次连接升级为websocket

//升级为websocket
ws, err := upgrader.Upgrade(w, r, nil)
if err != nil {
    log.Println(err)
    return
}
defer ws.Close()

//注意这里用的是缓存方式，有可能出现缓存溢出的现象
var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true // 不推荐在生产中使用
	},
}

2.7 准备连接

// Request a pseudo terminal
// 创建虚拟终端
if err := session.RequestPty("xterm", 80, 40, ssh.TerminalModes{}); err != nil {
    log.Fatalf("request for pseudo terminal failed: %v", err)
    //todo 返回错误码
}

// Start a shell
// 创建连接程序
if err := session.Shell(); err != nil {
    log.Fatal(err)
    //todo 返回错误码
}

// 为了能够先输出机器的标识符
// 如果不理解下面的代码可以先删除掉，你可以看到程序的第一次输出会缺少 机器标识符
// 我们将数据写到我们之前获取的管道对象
commands := []string{" "}
for _, cmd := range commands {
    _, err = stdinPipe.Write([]byte(cmd))
    if err != nil {
        log.Fatal(err)
    }
}

2.8绑定缓冲区

我们将远程终端响应的数据先中转到后端的缓存中，不是直接发向前端（因为可能在后端有些其余的处理）

我们将ssh输出管道对象绑定到我们自定义的缓冲区中，注意读写冲突问题。

/// Custom buffers
var stdoutBuf model.OutputDataBuffer
// Capture stdout and stderr
go io.Copy(&stdoutBuf, stdoutPipe)
go io.Copy(&stdoutBuf, stderrPipe)

//将虚拟机的初始化信息输出到ws上
//延时（如果不延时会让第一条数据被迅速覆盖）
time.Sleep(50 * time.Millisecond)
stdoutBuf.Flush(ws, connctID)

下面是缓存的代码

package model

import (
"bytes"
"github.com/gorilla/websocket"
"io"
"log"
"sync"
)

//本文件定义了与虚拟机进行交互的相关对象

// 线程安全的缓冲区
type OutputDataBuffer struct {
buffer bytes.Buffer
mu     sync.Mutex
}

func (w *OutputDataBuffer) Write(p []byte) (n int, err error) {
w.mu.Lock()
defer w.mu.Unlock()
return w.buffer.Write(p)
}

func (w *OutputDataBuffer) Flush(ws *websocket.Conn, connectID string) error {
w.mu.Lock()
defer w.mu.Unlock()
if w.buffer.Len() != 0 {
    //err := ws.WriteJSON(map[string]string{"Code": "200", "connectID": connectID, "cmdResponseData": w.buffer.String(), "message": ""})
    err := ws.WriteJSON(map[string]string{"Code": "200", "connectID": connectID, "resData": w.buffer.String(), "message": ""})
    if err != nil {
        return err
    }
    w.buffer.Reset()
}
log.Printf("缓冲区数据已经刷新到WebSocket连接中")
return nil
}

func (w *OutputDataBuffer) WriteToVM(stdinPipe io.WriteCloser) {
w.mu.Lock()
defer w.mu.Unlock()
stdinPipe.Write(w.buffer.Bytes())
w.buffer.Reset()
}

2.9 接收前端的数据与向前端响应数据

// 接受前端的数据

var inputBuffer model.OutputDataBuffer

for {
    _, message, err := ws.ReadMessage()
    if err != nil {
        log.Println("read:", err)
        break
    }

    var cmdData model.ShellRequest
    err = json.Unmarshal(message, &cmdData)
    if err != nil {
        log.Println("前端请求对象解析错误:", err)
        continue // 解析失败时继续监听下一条消息
    }
    // 将接收到的消息放入缓冲区
    inputBuffer.Write([]byte(cmdData.Command))

    //将缓冲区的数据写入虚拟机
    inputBuffer.WriteToVM(stdinPipe)

    //延时
    time.Sleep(50 * time.Millisecond)

    // 虚拟机的数据写入websocket
    stdoutBuf.Flush(ws, connctID)

}