backup/ssh.go

package main

import (
	"io"
	"os"
	"time"

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

const SshDialTimeout = time.Duration(10 * time.Second)
const SshInactivityTimeout = time.Duration(time.Minute)

type Ssh struct {
	name    string
	signer  ssh.Signer
	config  *ssh.ClientConfig
	client  *ssh.Client
	session *ssh.Session
	in      io.WriteCloser
	out     io.Reader
	err     io.Reader
}

func NewSsh(name, addr, user, key string) (*Ssh, error) {
	log.WithFields(log.Fields{"name": name, "addr": addr, "user": user, "key": key}).Debugf("starting")
	defer log.WithFields(log.Fields{"name": name, "addr": addr, "user": user, "key": key}).Debugf("done")

	s := &Ssh{
		name: name,
	}

	k, err := os.ReadFile(key)
	if err != nil {
		log.WithFields(log.Fields{"name": name, "addr": addr, "user": user, "key": key, "call": "os.ReadFile", "error": err}).Errorf("")
		return s, err
	}

	parsedKey, err := ssh.ParseRawPrivateKey(k)
	if err != nil {
		log.WithFields(log.Fields{"name": name, "addr": addr, "user": user, "key": key, "call": "ssh.ParseRawPrivateKey", "error": err}).Errorf("")
		return s, err
	}

	s.signer, err = ssh.NewSignerFromKey(parsedKey)
	if err != nil {
		log.WithFields(log.Fields{"name": name, "addr": addr, "user": user, "key": key, "call": "ssh.NewSignerFromKey", "error": err}).Errorf("")
		return s, err
	}

	s.config = &ssh.ClientConfig{
		User: user,
		Auth: []ssh.AuthMethod{
			ssh.PublicKeys(s.signer),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         SshDialTimeout,
	}

	s.client, err = ssh.Dial("tcp", addr, s.config)
	if err != nil {
		log.WithFields(log.Fields{"name": name, "addr": addr, "user": user, "key": key, "call": "ssh.Dial", "error": err}).Errorf("")
		return s, err
	}

	return s, nil

}

func (s *Ssh) Close() error {
	log.WithFields(log.Fields{"name": s.name}).Debugf("starting")
	defer log.WithFields(log.Fields{"name": s.name}).Debugf("done")

	return s.client.Close()
}

func NewSshPool(name, addr, user, key string) (pool.Pool, error) {
	log.WithFields(log.Fields{"name": name, "addr": addr, "user": user, "key": key}).Debugf("starting")
	defer log.WithFields(log.Fields{"name": name, "addr": addr, "user": user, "key": key}).Debugf("done")

	//factory Specify the method to create the connection
	factory := func() (interface{}, error) { return NewSsh(name, addr, user, key) }

	// close Specify the method to close the connection
	close := func(v interface{}) error { return v.(*Ssh).Close() }

	// Create a connection pool: Initialize the number of connections to 0, the maximum idle connection is 2, and the maximum concurrent connection is 25
	poolConfig := &pool.Config{
		InitialCap: 0,
		MaxIdle:    2,
		MaxCap:     25,
		Factory:    factory,
		Close:      close,
		//Ping:       ping,
		//The maximum idle time of the connection, the connection exceeding this time will be closed, which can avoid the problem of automatic failure when connecting to EOF when idle
		IdleTimeout: SshInactivityTimeout,
	}

	return pool.NewChannelPool(poolConfig)
}

func (s *Ssh) Exec(cmd string) (string, error) {
	log.WithFields(log.Fields{"name": s.name, "cmd": cmd}).Debugf("starting")
	defer log.WithFields(log.Fields{"name": s.name, "cmd": cmd}).Debugf("done")

	if err := s.ExecPipe(cmd); err != nil {
		log.WithFields(log.Fields{"name": s.name, "cmd": cmd, "call": "ssh.ExecPipe", "error": err}).Errorf("")
		return "", err
	}
	defer s.session.Close()

	if err := s.session.Wait(); err != nil {
		log.WithFields(log.Fields{"name": s.name, "cmd": cmd, "call": "session.Setenv", "error": err}).Errorf("")
		return "", err
	}

	buf, err := io.ReadAll(s.out)
	if err != nil {
		log.WithFields(log.Fields{"name": s.name, "cmd": cmd, "call": "io.ReadAll", "error": err}).Errorf("")
		return "", err
	}

	return string(buf), nil

}

func (s *Ssh) ExecPipe(cmd string) error {
	log.WithFields(log.Fields{"name": s.name, "cmd": cmd}).Debugf("starting")
	defer log.WithFields(log.Fields{"name": s.name, "cmd": cmd}).Debugf("done")

	session, err := s.client.NewSession()
	if err != nil {
		log.WithFields(log.Fields{"name": s.name, "cmd": cmd, "call": "client.NewSession", "error": err}).Errorf("")
		return err
	}
	s.session = session

	if s.session.Setenv("TZ", cfg.Timezone); err != nil {
		log.WithFields(log.Fields{"name": s.name, "cmd": cmd, "call": "session.Setenv", "error": err}).Errorf("")
		return err
	}

	if s.in, err = s.session.StdinPipe(); err != nil {
		log.WithFields(log.Fields{"name": s.name, "cmd": cmd, "call": "session.StdinPipe", "error": err}).Errorf("")
		s.session.Close()
		return err
	}

	if s.out, err = s.session.StdoutPipe(); err != nil {
		log.WithFields(log.Fields{"name": s.name, "cmd": cmd, "call": "session.StdoutPipe", "error": err}).Errorf("")
		s.session.Close()
		return err
	}

	if s.err, err = s.session.StderrPipe(); err != nil {
		log.WithFields(log.Fields{"name": s.name, "cmd": cmd, "call": "session.StderrPipe", "error": err}).Errorf("")
		s.session.Close()
		return err
	}

	if err = s.session.Start(cmd); err != nil {
		log.WithFields(log.Fields{"name": s.name, "cmd": cmd, "call": "session.Start", "error": err}).Errorf("")
		s.session.Close()
		return err
	}

	return nil
}