fscan/plugins/services/neo4j/connector.go

package neo4j

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

	"github.com/neo4j/neo4j-go-driver/v4/neo4j"
	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/plugins/base"
)

// Neo4jConnection Neo4j连接包装器
type Neo4jConnection struct {
	driver neo4j.Driver
	target string
	info   string
	isAuth bool
}

// Neo4jConnector Neo4j连接器实现
type Neo4jConnector struct{}

// NewNeo4jConnector 创建Neo4j连接器
func NewNeo4jConnector() *Neo4jConnector {
	return &Neo4jConnector{}
}

// Connect 连接到Neo4j服务器（不进行认证）
func (c *Neo4jConnector) Connect(ctx context.Context, info *common.HostInfo) (interface{}, error) {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	timeout := time.Duration(common.Timeout) * time.Second
	
	// 构造Neo4j URL
	uri := fmt.Sprintf("bolt://%s:%s", info.Host, info.Ports)
	
	// 先尝试无认证连接
	driver, dbInfo, isAuth, err := c.createConnection(uri, "", "", timeout)
	if err != nil {
		// 检查是否是Neo4j相关错误
		if c.isNeo4jError(err) {
			// 即使连接失败，但可以识别为Neo4j服务
			return &Neo4jConnection{
				driver: nil,
				target: target,
				info:   "Neo4j Graph Database (Service Detected)",
				isAuth: true,
			}, nil
		}
		return nil, err
	}
	
	return &Neo4jConnection{
		driver: driver,
		target: target,
		info:   dbInfo,
		isAuth: isAuth,
	}, nil
}

// Authenticate 使用凭据进行认证
func (c *Neo4jConnector) Authenticate(ctx context.Context, conn interface{}, cred *base.Credential) error {
	neo4jConn, ok := conn.(*Neo4jConnection)
	if !ok {
		return fmt.Errorf("invalid connection type")
	}
	
	// 解析目标地址
	parts := strings.Split(neo4jConn.target, ":")
	if len(parts) != 2 {
		return fmt.Errorf("invalid target format")
	}
	
	host := parts[0]
	port := parts[1]
	uri := fmt.Sprintf("bolt://%s:%s", host, port)
	timeout := time.Duration(common.Timeout) * time.Second
	
	// 使用提供的凭据创建新连接
	driver, info, isAuth, err := c.createConnection(uri, cred.Username, cred.Password, timeout)
	if err != nil {
		return err
	}
	
	// 更新连接信息
	if neo4jConn.driver != nil {
		neo4jConn.driver.Close()
	}
	neo4jConn.driver = driver
	neo4jConn.info = info
	neo4jConn.isAuth = isAuth
	
	return nil
}

// Close 关闭连接
func (c *Neo4jConnector) Close(conn interface{}) error {
	if neo4jConn, ok := conn.(*Neo4jConnection); ok && neo4jConn.driver != nil {
		return neo4jConn.driver.Close()
	}
	return nil
}

// createConnection 创建Neo4j连接
func (c *Neo4jConnector) createConnection(uri, username, password string, timeout time.Duration) (neo4j.Driver, string, bool, error) {
	// 配置驱动选项
	config := func(c *neo4j.Config) {
		c.SocketConnectTimeout = timeout
		c.ConnectionAcquisitionTimeout = timeout
		// Neo4j驱动默认不支持代理，这里暂不处理Socks代理
	}
	
	var driver neo4j.Driver
	var err error
	isAuth := true
	
	// 尝试建立连接
	if username != "" || password != "" {
		// 有认证信息时使用认证
		driver, err = neo4j.NewDriver(uri, neo4j.BasicAuth(username, password, ""), config)
	} else {
		// 无认证时使用NoAuth
		driver, err = neo4j.NewDriver(uri, neo4j.NoAuth(), config)
		isAuth = false
	}
	
	if err != nil {
		return nil, "", isAuth, err
	}
	
	// 测试连接有效性
	err = driver.VerifyConnectivity()
	if err != nil {
		driver.Close()
		return nil, "", isAuth, err
	}
	
	// 获取数据库信息
	info := c.getDatabaseInfo(driver)
	
	return driver, info, isAuth, nil
}

// getDatabaseInfo 获取Neo4j数据库信息
func (c *Neo4jConnector) getDatabaseInfo(driver neo4j.Driver) string {
	session := driver.NewSession(neo4j.SessionConfig{AccessMode: neo4j.AccessModeRead})
	defer session.Close()
	
	// 尝试获取版本信息
	result, err := session.Run("CALL dbms.components() YIELD name, versions, edition RETURN name, versions[0] as version, edition", nil)
	if err != nil {
		return "Neo4j Graph Database"
	}
	
	if result.Next() {
		record := result.Record()
		if name, ok := record.Get("name"); ok {
			if version, ok := record.Get("version"); ok {
				if edition, ok := record.Get("edition"); ok {
					return fmt.Sprintf("%s %s (%s)", name, version, edition)
				}
				return fmt.Sprintf("%s %s", name, version)
			}
		}
	}
	
	return "Neo4j Graph Database"
}

// isNeo4jError 检查是否是Neo4j相关错误
func (c *Neo4jConnector) isNeo4jError(err error) bool {
	if err == nil {
		return false
	}
	
	errorStr := strings.ToLower(err.Error())
	neo4jErrorIndicators := []string{
		"neo4j",
		"bolt",
		"authentication failed",
		"credentials",
		"unauthorized",
		"connection refused",
		"graph database",
		"cypher",
	}
	
	for _, indicator := range neo4jErrorIndicators {
		if strings.Contains(errorStr, indicator) {
			return true
		}
	}
	
	return false
}