be.ems/src/framework/utils/telnet/telnet.go

package telnet

import (
	"bytes"
	"fmt"
	"net"
	"strings"
	"time"
)

// ConnTelnet 连接telnet对象
type ConnTelnet struct {
	User     string `json:"user"`     // 主机用户名
	Addr     string `json:"addr"`     // 主机地址
	Port     int64  `json:"port"`     // telnet端口
	Password string `json:"password"` // 认证密码

	DialTimeOut time.Duration `json:"dialTimeOut"` // 连接超时断开

	Client     *net.Conn `json:"client"`
	LastResult string    `json:"lastResult"` // 记最后一次执行命令的结果
}

// NewClient 创建Telnet客户端
func (c *ConnTelnet) NewClient() (*ConnTelnet, error) {
	// IPV6地址协议
	proto := "tcp"
	if strings.Contains(c.Addr, ":") {
		proto = "tcp6"
		c.Addr = fmt.Sprintf("[%s]", c.Addr)
	}
	addr := fmt.Sprintf("%s:%d", c.Addr, c.Port)

	// 默认等待5s
	if c.DialTimeOut == 0 {
		c.DialTimeOut = 5 * time.Second
	}

	// 连接到 Telnet 服务器
	client, err := net.DialTimeout(proto, addr, c.DialTimeOut)
	if err != nil {
		return nil, err
	}

	// 需要确保接收方理解并正确处理发送窗口大小设置命令
	client.Write([]byte{255, 251, 31})                           // 发送窗口大小选项
	client.Write([]byte{255, 250, 31, 0, 128, 0, 120, 255, 240}) // 发送窗口行和列的大小

	// 进行登录
	time.Sleep(100 * time.Millisecond)
	client.Write([]byte(c.User + "\r\n"))
	time.Sleep(100 * time.Millisecond)
	client.Write([]byte(c.Password + "\r\n"))
	// fmt.Fprintln(client, c.User)
	// fmt.Fprintln(client, c.Password)
	c.Client = &client

	// 排空连接登录的信息
	c.RunCMD("")
	return c, nil
}

// Close 关闭当前Telnet客户端
func (c *ConnTelnet) Close() {
	if c.Client != nil {
		(*c.Client).Close()
	}
}

// RunCMD 执行单次命令
func (c *ConnTelnet) RunCMD(cmd string) (string, error) {
	if c.Client == nil {
		return "", fmt.Errorf("telnet client not connected")
	}
	conn := *c.Client
	var buf bytes.Buffer
	tmp := make([]byte, 1024)

	// 写入命令
	if cmd != "" {
		if _, err := conn.Write([]byte(cmd)); err != nil {
			return "", err
		}
	}

	// 读取命令消息
	for {
		// 设置读取超时时间为1000毫秒
		conn.SetReadDeadline(time.Now().Add(1000 * time.Millisecond))
		n, err := conn.Read(tmp)
		if err != nil {
			// 判断是否是超时错误
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				break
			}
			break
		}
		if n == 0 {
			break
		}
		buf.Write(tmp[:n])
	}
	defer buf.Reset()

	c.LastResult = buf.String()
	return c.LastResult, nil
}

// NewClient 创建Telnet客户端会话对象
func (c *ConnTelnet) NewClientSession(cols, rows uint8) (*TelnetClientSession, error) {
	if c.Client == nil {
		return nil, fmt.Errorf("telnet client not connected")
	}
	return &TelnetClientSession{
		Client: *c.Client,
	}, nil
}

// TelnetClientSession Telnet客户端会话对象
type TelnetClientSession struct {
	Client net.Conn
}

// Close 关闭会话
func (s *TelnetClientSession) Close() {
	if s.Client != nil {
		s.Client.Close()
	}
}

// Write 写入命令 不带回车(\n)也会执行根据客户端情况
func (s *TelnetClientSession) Write(cmd string) (int, error) {
	if s.Client == nil {
		return 0, fmt.Errorf("client is nil to content write failed")
	}
	return s.Client.Write([]byte(cmd))
}

// Read 读取结果 等待一会才有结果
func (s *TelnetClientSession) Read() []byte {
	if s.Client == nil {
		return []byte{}
	}

	buf := make([]byte, 1024)
	// 设置读取超时时间为100毫秒
	s.Client.SetReadDeadline(time.Now().Add(100 * time.Millisecond))
	_, err := s.Client.Read(buf)
	if err != nil {
		return []byte{}
	}
	return buf
}

// CombinedOutput 发送命令带结果返回
func (s *TelnetClientSession) CombinedOutput(cmd string) (string, error) {
	n, err := s.Write(cmd)
	if n == 0 || err != nil {
		return "", err
	}

	var buf bytes.Buffer
	tmp := make([]byte, 1024)
	for {
		// 设置读取超时时间为1000毫秒
		s.Client.SetReadDeadline(time.Now().Add(1000 * time.Millisecond))
		n, err := s.Client.Read(tmp)
		if err != nil {
			// 判断是否是超时错误
			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
				break
			}
			break
		}
		if n == 0 {
			break
		}
		buf.Write(tmp[:n])
	}
	defer buf.Reset()

	return buf.String(), nil
}