fscan/Plugins/services/redis/connector.go

package redis

import (
	"context"
	"fmt"
	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/plugins/base"
	"io"
	"net"
	"strings"
	"time"
)

// RedisConnection Redis连接包装
type RedisConnection struct {
	conn net.Conn
	authenticated bool
	config *RedisConfig
}

// RedisConfig Redis配置信息
type RedisConfig struct {
	DBFilename string
	Dir        string
}

// RedisConnector Redis连接器实现
type RedisConnector struct {
	timeout time.Duration
}

// NewRedisConnector 创建Redis连接器
func NewRedisConnector() *RedisConnector {
	return &RedisConnector{
		timeout: time.Duration(common.Timeout) * time.Second,
	}
}

// Connect 连接到Redis服务
func (c *RedisConnector) Connect(ctx context.Context, info *common.HostInfo) (interface{}, error) {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	
	// 建立TCP连接
	conn, err := common.WrapperTcpWithTimeout("tcp", target, c.timeout)
	if err != nil {
		return nil, fmt.Errorf("连接失败: %v", err)
	}
	
	// 创建Redis连接包装
	redisConn := &RedisConnection{
		conn: conn,
		authenticated: false,
		config: &RedisConfig{},
	}
	
	return redisConn, nil
}

// Authenticate 认证
func (c *RedisConnector) Authenticate(ctx context.Context, conn interface{}, cred *base.Credential) error {
	redisConn, ok := conn.(*RedisConnection)
	if !ok {
		return fmt.Errorf("无效的连接类型")
	}
	
	// 如果没有密码，先检查未授权访问
	if cred == nil || cred.Password == "" {
		return c.checkUnauthorizedAccess(redisConn)
	}
	
	// 有密码的情况下进行认证
	return c.authenticateWithPassword(redisConn, cred.Password)
}

// Close 关闭连接
func (c *RedisConnector) Close(conn interface{}) error {
	if redisConn, ok := conn.(*RedisConnection); ok {
		if redisConn.conn != nil {
			return redisConn.conn.Close()
		}
	}
	return nil
}

// checkUnauthorizedAccess 检查未授权访问
func (c *RedisConnector) checkUnauthorizedAccess(conn *RedisConnection) error {
	// 发送INFO命令测试
	if err := c.sendCommand(conn, "INFO"); err != nil {
		return fmt.Errorf("发送INFO命令失败: %v", err)
	}
	
	// 读取响应
	response, err := c.readResponse(conn)
	if err != nil {
		return fmt.Errorf("读取响应失败: %v", err)
	}
	
	// 检查是否包含Redis版本信息
	if !strings.Contains(response, "redis_version") {
		return fmt.Errorf("未发现Redis未授权访问")
	}
	
	// 获取配置信息
	if err := c.getConfig(conn); err != nil {
		common.LogDebug(fmt.Sprintf("获取Redis配置失败: %v", err))
	}
	
	conn.authenticated = true
	return nil
}

// authenticateWithPassword 使用密码认证
func (c *RedisConnector) authenticateWithPassword(conn *RedisConnection, password string) error {
	// 发送AUTH命令
	authCmd := fmt.Sprintf("AUTH %s", password)
	if err := c.sendCommand(conn, authCmd); err != nil {
		return fmt.Errorf("发送AUTH命令失败: %v", err)
	}
	
	// 读取响应
	response, err := c.readResponse(conn)
	if err != nil {
		return fmt.Errorf("读取响应失败: %v", err)
	}
	
	// 检查认证结果
	if !strings.Contains(response, "+OK") {
		return fmt.Errorf("认证失败")
	}
	
	// 获取配置信息
	if err := c.getConfig(conn); err != nil {
		common.LogDebug(fmt.Sprintf("获取Redis配置失败: %v", err))
	}
	
	conn.authenticated = true
	return nil
}

// sendCommand 发送Redis命令
func (c *RedisConnector) sendCommand(conn *RedisConnection, command string) error {
	// 设置写超时
	conn.conn.SetWriteDeadline(time.Now().Add(c.timeout))
	
	// 发送命令（添加CRLF）
	_, err := conn.conn.Write([]byte(command + "\r\n"))
	return err
}

// readResponse 读取Redis响应
func (c *RedisConnector) readResponse(conn *RedisConnection) (string, error) {
	// 设置读超时
	conn.conn.SetReadDeadline(time.Now().Add(c.timeout))
	
	// 读取所有数据
	data, err := io.ReadAll(conn.conn)
	if len(data) > 0 {
		// 如果读到数据，忽略EOF错误
		err = nil
	}
	
	return string(data), err
}

// getConfig 获取Redis配置
func (c *RedisConnector) getConfig(conn *RedisConnection) error {
	// 获取数据库文件名
	if err := c.sendCommand(conn, "CONFIG GET dbfilename"); err != nil {
		return err
	}
	
	response, err := c.readResponse(conn)
	if err != nil {
		return err
	}
	
	// 解析响应
	lines := strings.Split(response, "\r\n")
	if len(lines) > 2 {
		conn.config.DBFilename = lines[len(lines)-2]
	}
	
	// 获取数据库目录
	if err := c.sendCommand(conn, "CONFIG GET dir"); err != nil {
		return err
	}
	
	response, err = c.readResponse(conn)
	if err != nil {
		return err
	}
	
	// 解析响应
	lines = strings.Split(response, "\r\n")
	if len(lines) > 2 {
		conn.config.Dir = lines[len(lines)-2]
	}
	
	return nil
}

// =============================================================================
// Redis操作辅助函数
// =============================================================================

// ExecuteCommand 执行Redis命令
func (c *RedisConnector) ExecuteCommand(conn *RedisConnection, command string) (string, error) {
	if !conn.authenticated {
		return "", fmt.Errorf("连接未认证")
	}
	
	if err := c.sendCommand(conn, command); err != nil {
		return "", err
	}
	
	return c.readResponse(conn)
}

// SetConfig 设置Redis配置
func (c *RedisConnector) SetConfig(conn *RedisConnection, key, value string) error {
	if !conn.authenticated {
		return fmt.Errorf("连接未认证")
	}
	
	command := fmt.Sprintf("CONFIG SET %s %s", key, value)
	if err := c.sendCommand(conn, command); err != nil {
		return err
	}
	
	response, err := c.readResponse(conn)
	if err != nil {
		return err
	}
	
	if !strings.Contains(response, "OK") {
		return fmt.Errorf("设置配置失败: %s", response)
	}
	
	return nil
}

// SetKey 设置Redis键值
func (c *RedisConnector) SetKey(conn *RedisConnection, key, value string) error {
	if !conn.authenticated {
		return fmt.Errorf("连接未认证")
	}
	
	command := fmt.Sprintf("SET %s \"%s\"", key, value)
	if err := c.sendCommand(conn, command); err != nil {
		return err
	}
	
	response, err := c.readResponse(conn)
	if err != nil {
		return err
	}
	
	if !strings.Contains(response, "OK") {
		return fmt.Errorf("设置键值失败: %s", response)
	}
	
	return nil
}

// Save 保存Redis数据
func (c *RedisConnector) Save(conn *RedisConnection) error {
	if !conn.authenticated {
		return fmt.Errorf("连接未认证")
	}
	
	if err := c.sendCommand(conn, "SAVE"); err != nil {
		return err
	}
	
	response, err := c.readResponse(conn)
	if err != nil {
		return err
	}
	
	if !strings.Contains(response, "OK") {
		return fmt.Errorf("保存数据失败: %s", response)
	}
	
	return nil
}

// RestoreConfig 恢复Redis配置
func (c *RedisConnector) RestoreConfig(conn *RedisConnection, originalConfig *RedisConfig) error {
	if originalConfig.DBFilename != "" {
		if err := c.SetConfig(conn, "dbfilename", originalConfig.DBFilename); err != nil {
			return fmt.Errorf("恢复dbfilename失败: %v", err)
		}
	}
	
	if originalConfig.Dir != "" {
		if err := c.SetConfig(conn, "dir", originalConfig.Dir); err != nil {
			return fmt.Errorf("恢复dir失败: %v", err)
		}
	}
	
	return nil
}