s_cherry_e/sr/tool/cleandb/cleandb.go

package main

import (
	"context"
	"fmt"
	"net"
	"time"

	"github.com/go-redis/redis/v8"
	"go.mongodb.org/mongo-driver/v2/mongo"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
	"golang.org/x/crypto/ssh"
)

// SSHConfig SSH 连接配置
type SSHConfig struct {
	Host     string // SSH 转发服务器地址, e.g. "1.2.3.4"
	Port     int    // SSH 端口, e.g. 22
	User     string // SSH 用户名
	Password string // SSH 密码 (若使用密钥则留空)
}

// DBConfig 数据库连接配置
type DBConfig struct {
	Host       string // 数据库地址 (通常是 127.0.0.1)
	Port       int    // 数据库端口
	User       string // 数据库用户名
	Password   string // 数据库密码
	AuthSource string // 认证数据库 (如 "admin")
}

// SSHTunnel SSH 隧道包装
type SSHTunnel struct {
	client *ssh.Client
}

// Close 关闭 SSH 连接
func (s *SSHTunnel) Close() error {
	if s.client != nil {
		return s.client.Close()
	}
	return nil
}

// createSSHTunnel 创建 SSH 客户端并返回拨号器
func createSSHTunnel(sshConf SSHConfig) (*SSHTunnel, error) {
	config := &ssh.ClientConfig{
		User: sshConf.User,
		Auth: []ssh.AuthMethod{
			ssh.Password(sshConf.Password),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         10 * time.Second,
	}

	addr := fmt.Sprintf("%s:%d", sshConf.Host, sshConf.Port)
	client, err := ssh.Dial("tcp", addr, config)
	if err != nil {
		return nil, fmt.Errorf("failed to dial ssh: %w", err)
	}

	return &SSHTunnel{client: client}, nil
}

// sshConn 包装 net.Conn 以支持 MongoDB 和 Redis 驱动所需的 Deadline 方法
// 因为 SSH 通道 (channel) 默认不支持 SetDeadline
type sshConn struct {
	net.Conn
}

func (c *sshConn) SetDeadline(t time.Time) error      { return nil }
func (c *sshConn) SetReadDeadline(t time.Time) error  { return nil }
func (c *sshConn) SetWriteDeadline(t time.Time) error { return nil }

// mongoDialer 实现 options.ContextDialer 接口
type mongoDialer struct {
	tunnel *SSHTunnel
}

func (m *mongoDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
	conn, err := m.tunnel.client.Dial(network, address)
	if err != nil {
		return nil, err
	}
	return &sshConn{conn}, nil
}

// ConnectMongoViaSSH 通过 SSH 隧道连接 MongoDB (若 sshConf.Host 为空则直连)
func ConnectMongoViaSSH(sshConf SSHConfig, dbConf DBConfig) (*mongo.Client, *SSHTunnel, error) {
	var tunnel *SSHTunnel
	var dialer options.ContextDialer
	var err error

	if sshConf.Host != "" {
		tunnel, err = createSSHTunnel(sshConf)
		if err != nil {
			return nil, nil, err
		}
		dialer = &mongoDialer{tunnel: tunnel}
	}

	mongoAddr := fmt.Sprintf("%s:%d", dbConf.Host, dbConf.Port)
	uri := fmt.Sprintf("mongodb://%s", mongoAddr)

	clientOptions := options.Client().ApplyURI(uri)
	if dialer != nil {
		clientOptions.SetDialer(dialer)
	}

	if dbConf.User != "" {
		if dbConf.AuthSource == "" {
			dbConf.AuthSource = "admin"
		}

		clientOptions.SetAuth(options.Credential{
			Username:   dbConf.User,
			Password:   dbConf.Password,
			AuthSource: dbConf.AuthSource,
		})
	}

	client, err := mongo.Connect(clientOptions)
	if err != nil {
		if tunnel != nil {
			tunnel.Close()
		}
		return nil, nil, fmt.Errorf("failed to connect to mongo: %w", err)
	}

	return client, tunnel, nil
}

// ConnectRedisViaSSH 通过 SSH 隧道连接 Redis (若 sshConf.Host 为空则直连)
func ConnectRedisViaSSH(sshConf SSHConfig, dbConf DBConfig) (*redis.Client, *SSHTunnel, error) {
	var tunnel *SSHTunnel
	var err error

	redisAddr := fmt.Sprintf("%s:%d", dbConf.Host, dbConf.Port)
	redisOptions := &redis.Options{
		Addr:     redisAddr,
		Username: dbConf.User,
		Password: dbConf.Password,
		DB:       0, // 默认 DB
	}

	if sshConf.Host != "" {
		tunnel, err = createSSHTunnel(sshConf)
		if err != nil {
			return nil, nil, err
		}
		redisOptions.Dialer = func(ctx context.Context, network, addr string) (net.Conn, error) {
			conn, err := tunnel.client.Dial(network, addr)
			if err != nil {
				return nil, err
			}
			return &sshConn{conn}, nil
		}
	}

	rdb := redis.NewClient(redisOptions)

	// 验证连接
	if err := rdb.Ping(context.Background()).Err(); err != nil {
		rdb.Close()
		if tunnel != nil {
			tunnel.Close()
		}
		return nil, nil, fmt.Errorf("failed to ping redis: %w", err)
	}

	return rdb, tunnel, nil
}