golang ssh库封装

概述

在日常运维工作中需要写一些小工具实现远程知行命令并返回执行结果，python 有 paramiko 库，golang 原生的 ssh 库使用有点小复杂，这里记录一下，后面使用就可以直接 copy 了。

上代码

package ssh

import (
    "bufio"
    "fmt"
    "io/ioutil"
    "net"
    "strings"

    log "github.com/sirupsen/logrus"
    "golang.org/x/crypto/ssh"
)

type Client struct {
    *ssh.Client
}

func PublicKeyFile(file string) ssh.AuthMethod {
    buffer, err := ioutil.ReadFile(file)
    if err != nil {
        log.Fatalln(err)
    }

    key, err := ssh.ParsePrivateKey(buffer)
    if err != nil {
        log.Fatalln(err)
    }
    return ssh.PublicKeys(key)
}

func NewClient(user, keyFile, addr string) (*Client, error) {
    config := &ssh.ClientConfig{
        User: user,
        Auth: []ssh.AuthMethod{PublicKeyFile(keyFile)},
        HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
            return nil
        },
    }

    client, err := ssh.Dial("tcp", addr, config)
    if err != nil {
        return nil, fmt.Errorf("could not establish ssh connection: %w", err)
    }

    return &Client{client}, nil
}

func (c *Client) Close() {
    c.Client.Close()
}

func (c *Client) ExecuteRemoteCommand(command string) error {
    session, err := c.NewSession()
    if err != nil {
        return fmt.Errorf("create new session failed: %w", err)
    }
    defer session.Close()

    modes := ssh.TerminalModes{
        ssh.ECHO:          1,
        ssh.TTY_OP_ISPEED: 144,
        ssh.TTY_OP_OSPEED: 144,
    }

    if err := session.RequestPty("vt100", 80, 40, modes); err != nil {
        return fmt.Errorf("request for pseudo terminal failed: %w", err)
    }

    stdinBuf, _ := session.StdinPipe()
    defer stdinBuf.Close()

    stdErrBuf, _ := session.StderrPipe()
    errReader := bufio.NewScanner(stdErrBuf)

    stdOutBuf, _ := session.StdoutPipe()
    outReader := bufio.NewScanner(stdOutBuf)

    go func() {
        for outReader.Scan() || errReader.Scan() {
            if outReader.Text() != "" {
                line := strings.TrimSpace(outReader.Text())
                fmt.Println(">>>\t" + line)
            } else if errReader.Text() != "" {
                line := strings.TrimSpace(errReader.Text())
                fmt.Println(">>>ERR:\t" + line)
            }
        }
    }()

    if err := session.Start(command); err != nil {
        return fmt.Errorf("start command failed: %w", err)
    }

    if err := session.Wait(); err != nil {
        return fmt.Errorf("wait end of command execution failed: %w", err)
    }

    return nil
}

使用示例

func main() {
    sshConn, err := myssh.NewClient("username", "keyFile", "serverIp:port")
    if err != nil {
        log.Fatalf("connect to server failed: %v", err)
    }
    defer sshConn.Close()

    err = myssh.ExecuteRemoteCommand(sshConn, "longRunningCommand")
    if err != nil {
        log.Fatalf("run command failed: %v", err)
    }
}