88 lines
1.9 KiB
Go
88 lines
1.9 KiB
Go
package scheduler
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"net"
|
|
"os"
|
|
"time"
|
|
|
|
"golang.org/x/crypto/ssh"
|
|
)
|
|
|
|
// SSHClient SSH 连接客户端
|
|
type SSHClient struct {
|
|
client *ssh.Client
|
|
}
|
|
|
|
// Connect 建立 SSH 连接
|
|
func (s *SSHClient) Connect(host string, port int, user, password, privateKeySrc string, authType int) error {
|
|
addr := net.JoinHostPort(host, fmt.Sprintf("%d", port))
|
|
config := &ssh.ClientConfig{
|
|
User: user,
|
|
Timeout: 10 * time.Second,
|
|
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
|
|
}
|
|
|
|
if authType == 0 {
|
|
// 密码认证
|
|
config.Auth = []ssh.AuthMethod{ssh.Password(password)}
|
|
} else {
|
|
// 密钥认证
|
|
keyData, err := os.ReadFile(privateKeySrc)
|
|
if err != nil {
|
|
return fmt.Errorf("读取密钥文件失败: %w", err)
|
|
}
|
|
signer, err := ssh.ParsePrivateKey(keyData)
|
|
if err != nil {
|
|
return fmt.Errorf("解析密钥失败: %w", err)
|
|
}
|
|
config.Auth = []ssh.AuthMethod{ssh.PublicKeys(signer)}
|
|
}
|
|
|
|
client, err := ssh.Dial("tcp", addr, config)
|
|
if err != nil {
|
|
return fmt.Errorf("SSH连接失败: %w", err)
|
|
}
|
|
s.client = client
|
|
return nil
|
|
}
|
|
|
|
// Exec 执行远程命令
|
|
func (s *SSHClient) Exec(command string, timeoutSec int) (stdout, stderr string, err error) {
|
|
session, err := s.client.NewSession()
|
|
if err != nil {
|
|
return "", "", fmt.Errorf("创建SSH会话失败: %w", err)
|
|
}
|
|
defer session.Close()
|
|
|
|
var outBuf, errBuf bytes.Buffer
|
|
session.Stdout = &outBuf
|
|
session.Stderr = &errBuf
|
|
|
|
done := make(chan error, 1)
|
|
go func() {
|
|
done <- session.Run(command)
|
|
}()
|
|
|
|
timeout := time.Duration(timeoutSec) * time.Second
|
|
if timeoutSec <= 0 {
|
|
timeout = 24 * time.Hour
|
|
}
|
|
|
|
select {
|
|
case err := <-done:
|
|
return outBuf.String(), errBuf.String(), err
|
|
case <-time.After(timeout):
|
|
session.Signal(ssh.SIGKILL)
|
|
return "", fmt.Sprintf("任务执行超过%d秒，已强制终止", timeoutSec), fmt.Errorf("timeout")
|
|
}
|
|
}
|
|
|
|
// Close 关闭连接
|
|
func (s *SSHClient) Close() {
|
|
if s.client != nil {
|
|
s.client.Close()
|
|
}
|
|
}
|