基于Charmbracelet/Wish构建SSH服务器集成Cobra命令的实践指南-CSDN博客

本文链接：https://blog.csdn.net/gitblog_01049/article/details/148549299

基于Charmbracelet/Wish构建SSH服务器集成Cobra命令的实践指南

项目背景与概述

Charmbracelet/Wish是一个优雅的Go语言SSH服务器库，它简化了SSH服务器的创建过程，同时提供了丰富的中间件支持。本文将通过分析一个集成Cobra命令行工具的示例，展示如何利用Wish构建功能强大的SSH服务。

核心功能解析

这个示例展示了如何将Cobra命令行工具集成到SSH服务器中，使得用户通过SSH连接后可以直接执行预定义的命令。主要实现了以下功能：

创建基本的SSH服务器
集成Cobra命令行工具
实现字符串回显功能（支持反向输出）
完善的服务器生命周期管理

代码实现详解

1. 命令行工具定义

示例中首先定义了一个简单的Cobra命令cmd()，该命令实现字符串回显功能：

func cmd() *cobra.Command {
    var reverse bool
    cmd := &cobra.Command{
        Use:  "echo [string]",
        Args: cobra.ExactArgs(1),
        RunE: func(cmd *cobra.Command, args []string) error {
            s := args[0]
            if reverse {
                ss := make([]byte, 0, len(s))
                for i := len(s) - 1; i >= 0; i-- {
                    ss = append(ss, s[i])
                }
                s = string(ss)
            }
            cmd.Println(s)
            return nil
        },
    }
    cmd.PersistentFlags().BoolVarP(&reverse, "reverse", "r", false, "Reverse string on echo")
    return cmd
}

这个命令接受一个字符串参数，并通过-r标志决定是否反向输出字符串。这是一个典型的Cobra命令实现，展示了如何在SSH会话中使用命令行工具。

2. SSH服务器配置

主函数中配置了SSH服务器：

s, err := wish.NewServer(
    wish.WithAddress(net.JoinHostPort(host, port)),
    wish.WithHostKeyPath(".ssh/id_ed25519"),
    wish.WithMiddleware(
        func(next ssh.Handler) ssh.Handler {
            return func(sess ssh.Session) {
                rootCmd := cmd()
                rootCmd.SetArgs(sess.Command())
                rootCmd.SetIn(sess)
                rootCmd.SetOut(sess)
                rootCmd.SetErr(sess.Stderr())
                rootCmd.CompletionOptions.DisableDefaultCmd = true
                if err := rootCmd.Execute(); err != nil {
                    _ = sess.Exit(1)
                    return
                }
                next(sess)
            }
        },
        logging.Middleware(),
    ),
)

关键点在于中间件的实现，它将SSH会话与Cobra命令绑定：

为每个会话创建新的命令实例
将会话的命令参数传递给Cobra
将标准输入输出重定向到SSH会话
执行命令并处理错误

3. 服务器生命周期管理

示例展示了完整的服务器生命周期管理：

done := make(chan os.Signal, 1)
signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
log.Info("Starting SSH server", "host", host, "port", port)
go func() {
    if err = s.ListenAndServe(); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
        log.Error("Could not start server", "error", err)
        done <- nil
    }
}()

<-done
log.Info("Stopping SSH server")
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer func() { cancel() }()
if err := s.Shutdown(ctx); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
    log.Error("Could not stop server", "error", err)
}

这部分代码处理了：