fscan/Plugins/services/snmp/connector.go

package snmp

import (
	"context"
	"fmt"
	"strconv"
	"strings"
	"time"

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

// SNMPConnector SNMP连接器实现
type SNMPConnector struct {
	host string
	port int
}

// SNMPConnection SNMP连接结构
type SNMPConnection struct {
	client    *gosnmp.GoSNMP
	community string
	sysDesc   string
	info      string
}

// NewSNMPConnector 创建SNMP连接器
func NewSNMPConnector() *SNMPConnector {
	return &SNMPConnector{}
}

// Connect 建立SNMP连接
func (c *SNMPConnector) Connect(ctx context.Context, info *common.HostInfo) (interface{}, error) {
	// 解析端口
	port, err := strconv.Atoi(info.Ports)
	if err != nil {
		return nil, fmt.Errorf("无效的端口号: %s", info.Ports)
	}

	c.host = info.Host
	c.port = port

	timeout := time.Duration(common.Timeout) * time.Second

	// 结果通道
	type connResult struct {
		conn   *SNMPConnection
		err    error
		banner string
	}
	resultChan := make(chan connResult, 1)

	// 在协程中尝试连接
	go func() {
		// 尝试使用默认的public community进行连接
		client := &gosnmp.GoSNMP{
			Target:    info.Host,
			Port:      uint16(port),
			Community: "public",
			Version:   gosnmp.Version2c,
			Timeout:   timeout,
			Retries:   1,
		}

		err := client.Connect()
		if err != nil {
			select {
			case <-ctx.Done():
			case resultChan <- connResult{nil, err, ""}:
			}
			return
		}

		// 尝试获取系统描述信息
		oids := []string{"1.3.6.1.2.1.1.1.0"} // sysDescr OID
		result, err := client.Get(oids)
		
		var sysDesc string
		var banner string
		
		if err == nil && len(result.Variables) > 0 {
			if result.Variables[0].Type != gosnmp.NoSuchObject {
				switch v := result.Variables[0].Value.(type) {
				case []byte:
					sysDesc = strings.TrimSpace(string(v))
				case string:
					sysDesc = strings.TrimSpace(v)
				}
			}
		}

		if sysDesc != "" {
			banner = fmt.Sprintf("SNMP Service (Version: %s, System: %s)", client.Version.String(), sysDesc)
		} else {
			banner = fmt.Sprintf("SNMP Service (Version: %s)", client.Version.String())
		}

		// 创建连接对象
		snmpConn := &SNMPConnection{
			client:    client,
			community: "public",
			sysDesc:   sysDesc,
			info:      banner,
		}

		select {
		case <-ctx.Done():
			client.Conn.Close()
		case resultChan <- connResult{snmpConn, nil, banner}:
		}
	}()

	// 等待连接结果
	select {
	case result := <-resultChan:
		if result.err != nil {
			return nil, result.err
		}
		return result.conn, nil
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

// Authenticate 进行SNMP认证（通过community字符串）
func (c *SNMPConnector) Authenticate(ctx context.Context, conn interface{}, cred *base.Credential) error {
	snmpConn, ok := conn.(*SNMPConnection)
	if !ok {
		return fmt.Errorf("无效的SNMP连接类型")
	}

	// 对于SNMP，将用户名作为community字符串使用
	community := cred.Username
	if community == "" {
		community = cred.Password // 如果用户名为空，尝试使用密码作为community
	}
	
	timeout := time.Duration(common.Timeout) * time.Second

	// 结果通道
	type authResult struct {
		client  *gosnmp.GoSNMP
		sysDesc string
		err     error
	}
	resultChan := make(chan authResult, 1)

	// 在协程中尝试认证
	go func() {
		// 关闭旧连接
		if snmpConn.client != nil && snmpConn.client.Conn != nil {
			snmpConn.client.Conn.Close()
		}

		// 创建新的SNMP客户端
		client := &gosnmp.GoSNMP{
			Target:    c.host,
			Port:      uint16(c.port),
			Community: community,
			Version:   gosnmp.Version2c,
			Timeout:   timeout,
			Retries:   1,
		}

		err := client.Connect()
		if err != nil {
			select {
			case <-ctx.Done():
			case resultChan <- authResult{nil, "", err}:
			}
			return
		}

		// 尝试获取系统描述信息验证认证
		oids := []string{"1.3.6.1.2.1.1.1.0"} // sysDescr OID
		result, err := client.Get(oids)
		if err != nil {
			client.Conn.Close()
			select {
			case <-ctx.Done():
			case resultChan <- authResult{nil, "", err}:
			}
			return
		}

		var sysDesc string
		if len(result.Variables) > 0 && result.Variables[0].Type != gosnmp.NoSuchObject {
			switch v := result.Variables[0].Value.(type) {
			case []byte:
				sysDesc = strings.TrimSpace(string(v))
			case string:
				sysDesc = strings.TrimSpace(v)
			}
		}

		select {
		case <-ctx.Done():
			client.Conn.Close()
		case resultChan <- authResult{client, sysDesc, nil}:
		}
	}()

	// 等待认证结果
	select {
	case result := <-resultChan:
		if result.err != nil {
			return fmt.Errorf(i18n.GetText("snmp_auth_failed"), result.err)
		}

		// 更新连接信息
		snmpConn.client = result.client
		snmpConn.community = community
		snmpConn.sysDesc = result.sysDesc
		if result.sysDesc != "" {
			snmpConn.info = fmt.Sprintf("SNMP Service (Community: %s, System: %s)", community, result.sysDesc)
		} else {
			snmpConn.info = fmt.Sprintf("SNMP Service (Community: %s)", community)
		}
		
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

// Close 关闭SNMP连接
func (c *SNMPConnector) Close(conn interface{}) error {
	if snmpConn, ok := conn.(*SNMPConnection); ok {
		if snmpConn.client != nil && snmpConn.client.Conn != nil {
			snmpConn.client.Conn.Close()
		}
		return nil
	}
	return fmt.Errorf("无效的SNMP连接类型")
}

// GetConnectionInfo 获取连接信息
func (conn *SNMPConnection) GetConnectionInfo() map[string]interface{} {
	info := map[string]interface{}{
		"protocol":  "SNMP",
		"service":   "SNMP",
		"info":      conn.info,
		"community": conn.community,
	}

	if conn.sysDesc != "" {
		info["system"] = conn.sysDesc
	}

	return info
}

// IsAlive 检查连接是否仍然有效
func (conn *SNMPConnection) IsAlive() bool {
	if conn.client == nil || conn.client.Conn == nil {
		return false
	}

	// 简单的SNMP GET测试连接
	oids := []string{"1.3.6.1.2.1.1.1.0"}
	_, err := conn.client.Get(oids)
	return err == nil
}

// GetServerInfo 获取服务器信息
func (conn *SNMPConnection) GetServerInfo() string {
	return conn.info
}