fscan/Plugins/services/modbus/connector.go

package modbus

import (
	"context"
	"encoding/binary"
	"fmt"
	"net"
	"time"

	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/plugins/base"
)

// ModbusConnection Modbus连接包装器
type ModbusConnection struct {
	client   net.Conn
	target   string
	deviceID uint8
}

// ModbusConnector Modbus连接器实现
type ModbusConnector struct {
	host string
	port string
}

// NewModbusConnector 创建Modbus连接器
func NewModbusConnector() *ModbusConnector {
	return &ModbusConnector{}
}

// Connect 连接到Modbus服务
func (c *ModbusConnector) Connect(ctx context.Context, info *common.HostInfo) (interface{}, error) {
	c.host = info.Host
	c.port = info.Ports
	
	target := fmt.Sprintf("%s:%s", c.host, c.port)
	timeout := time.Duration(common.Timeout) * time.Second
	
	// 建立TCP连接
	conn, err := net.DialTimeout("tcp", target, timeout)
	if err != nil {
		return nil, fmt.Errorf("Modbus连接失败: %v", err)
	}
	
	return &ModbusConnection{
		client: conn,
		target: target,
	}, nil
}

// Authenticate 认证 - Modbus通常无认证，检查协议响应
func (c *ModbusConnector) Authenticate(ctx context.Context, conn interface{}, cred *base.Credential) error {
	modbusConn, ok := conn.(*ModbusConnection)
	if !ok {
		return fmt.Errorf("无效的连接类型")
	}
	
	// 在goroutine中进行操作，支持Context取消
	resultChan := make(chan error, 1)
	
	go func() {
		// 设置操作超时
		timeout := time.Duration(common.Timeout) * time.Second
		if err := modbusConn.client.SetDeadline(time.Now().Add(timeout)); err != nil {
			resultChan <- fmt.Errorf("设置超时失败: %v", err)
			return
		}
		
		// 构建Modbus TCP请求包 - 读取设备ID
		request := buildModbusReadCoilsRequest()
		
		// 发送请求
		if _, err := modbusConn.client.Write(request); err != nil {
			resultChan <- fmt.Errorf("发送Modbus请求失败: %v", err)
			return
		}
		
		// 读取响应
		response := make([]byte, 256)
		n, err := modbusConn.client.Read(response)
		if err != nil {
			resultChan <- fmt.Errorf("读取Modbus响应失败: %v", err)
			return
		}
		
		// 验证是否为有效Modbus响应
		if isValidModbusResponse(response[:n]) {
			// 提取设备ID
			if len(response) >= 7 {
				modbusConn.deviceID = response[6] // Unit ID
			}
			resultChan <- nil
			return
		}
		
		resultChan <- fmt.Errorf("非Modbus服务或访问被拒绝")
	}()
	
	// 等待操作结果或Context取消
	select {
	case err := <-resultChan:
		return err
	case <-ctx.Done():
		return fmt.Errorf("Modbus操作超时: %v", ctx.Err())
	}
}

// Close 关闭连接
func (c *ModbusConnector) Close(conn interface{}) error {
	if modbusConn, ok := conn.(*ModbusConnection); ok && modbusConn.client != nil {
		return modbusConn.client.Close()
	}
	return nil
}

// buildModbusReadCoilsRequest 构建Modbus TCP读取线圈请求包
func buildModbusReadCoilsRequest() []byte {
	request := make([]byte, 12)
	
	// Modbus TCP头部
	binary.BigEndian.PutUint16(request[0:], 0x0001) // 事务标识符
	binary.BigEndian.PutUint16(request[2:], 0x0000) // 协议标识符
	binary.BigEndian.PutUint16(request[4:], 0x0006) // 长度
	request[6] = 0x01                               // 单元标识符
	
	// Modbus 功能码和数据
	request[7] = 0x01                                // 功能码: Read Coils
	binary.BigEndian.PutUint16(request[8:], 0x0000)  // 起始地址
	binary.BigEndian.PutUint16(request[10:], 0x0001) // 读取数量
	
	return request
}

// isValidModbusResponse 验证Modbus响应是否有效
func isValidModbusResponse(response []byte) bool {
	if len(response) < 8 {
		return false
	}
	
	// 检查协议标识符 (应为0)
	protocolID := binary.BigEndian.Uint16(response[2:])
	if protocolID != 0 {
		return false
	}
	
	// 检查功能码是否为错误响应
	funcCode := response[7]
	if funcCode >= 0x80 { // 错误响应的功能码都大于等于0x80
		return false
	}
	
	// 基本长度检查
	length := binary.BigEndian.Uint16(response[4:])
	if int(length) != len(response)-6 {
		return false
	}
	
	return true
}