aiops/backend/pkg/winsmb/client.go

package winsmb

import (
	"context"
	"fmt"
	"io"
	"net"
	"os"
	"path"
	"strings"
	"time"

	"github.com/hirochachacha/go-smb2"
)

// 虚拟 Download 盘在 Windows 上的相对路径（挂载 C$ 共享）// Logan.Li
const (
	AdminShare       = "C$"
	DownloadRelDir   = `Users\Public\AIOpsShare\Download`
	DownloadFullPath = `C:\Users\Public\AIOpsShare\Download`
)

// Auth SMB 管理员凭据 // Logan.Li
type Auth struct {
	Host     string
	Domain   string
	Username string
	Password string
}

// FileEntry 目录项 // Logan.Li
type FileEntry struct {
	Name    string
	Size    int64
	IsDir   bool
	ModTime time.Time
}

func sanitizeFileName(name string) string {
	base := path.Base(strings.ReplaceAll(name, `\`, "/"))
	base = strings.TrimSpace(base)
	if base == "" || base == "." || base == ".." {
		return "upload.bin"
	}
	return base
}

func (a Auth) dial() (*smb2.Share, func(), error) {
	if strings.TrimSpace(a.Host) == "" {
		return nil, nil, fmt.Errorf("SMB 主机为空")
	}
	d := &smb2.Dialer{
		Initiator: &smb2.NTLMInitiator{
			User:     a.Username,
			Password: a.Password,
			Domain:   a.Domain,
		},
	}
	tcpConn, err := net.DialTimeout("tcp", net.JoinHostPort(a.Host, "445"), 15*time.Second)
	if err != nil {
		return nil, nil, fmt.Errorf("SMB 连接失败（请确认 445 端口可达且账号有管理员权限）: %w", err)
	}
	session, err := d.DialContext(context.Background(), tcpConn)
	if err != nil {
		tcpConn.Close()
		return nil, nil, fmt.Errorf("SMB 认证失败: %w", err)
	}
	share, err := session.Mount(AdminShare)
	if err != nil {
		_ = session.Logoff()
		return nil, nil, fmt.Errorf("挂载 %s 失败: %w", AdminShare, err)
	}
	cleanup := func() {
		_ = share.Umount()
		_ = session.Logoff()
	}
	return share, cleanup, nil
}

// EnsureDownloadDir 创建虚拟 Download 目录 // Logan.Li
func EnsureDownloadDir(auth Auth) error {
	share, closeFn, err := auth.dial()
	if err != nil {
		return err
	}
	defer closeFn()
	cur := ""
	for _, part := range strings.Split(DownloadRelDir, `\`) {
		if part == "" {
			continue
		}
		if cur == "" {
			cur = part
		} else {
			cur = cur + `\` + part
		}
		_ = share.Mkdir(cur, 0755)
	}
	return nil
}

func remoteFilePath(fileName string) string {
	return DownloadRelDir + `\` + sanitizeFileName(fileName)
}

// Upload 上传文件到虚拟 Download 盘 // Logan.Li
func Upload(auth Auth, fileName string, data []byte) error {
	if err := EnsureDownloadDir(auth); err != nil {
		return err
	}
	share, closeFn, err := auth.dial()
	if err != nil {
		return err
	}
	defer closeFn()
	target := remoteFilePath(fileName)
	f, err := share.Create(target)
	if err != nil {
		return fmt.Errorf("创建远端文件失败: %w", err)
	}
	defer f.Close()
	if _, err := f.Write(data); err != nil {
		return fmt.Errorf("写入远端文件失败: %w", err)
	}
	return nil
}

// Download 从虚拟 Download 盘读取文件 // Logan.Li
func Download(auth Auth, fileName string) ([]byte, error) {
	share, closeFn, err := auth.dial()
	if err != nil {
		return nil, err
	}
	defer closeFn()
	target := remoteFilePath(fileName)
	f, err := share.Open(target)
	if err != nil {
		if os.IsNotExist(err) {
			return nil, fmt.Errorf("文件不存在: %s", sanitizeFileName(fileName))
		}
		return nil, fmt.Errorf("打开远端文件失败: %w", err)
	}
	defer f.Close()
	return io.ReadAll(f)
}

// ListDownload 列出虚拟 Download 盘文件 // Logan.Li
func ListDownload(auth Auth) ([]FileEntry, error) {
	if err := EnsureDownloadDir(auth); err != nil {
		return nil, err
	}
	share, closeFn, err := auth.dial()
	if err != nil {
		return nil, err
	}
	defer closeFn()
	entries, err := share.ReadDir(DownloadRelDir)
	if err != nil {
		return nil, fmt.Errorf("读取目录失败: %w", err)
	}
	out := make([]FileEntry, 0, len(entries))
	for _, ent := range entries {
		out = append(out, FileEntry{
			Name:    ent.Name(),
			Size:    ent.Size(),
			IsDir:   ent.IsDir(),
			ModTime: ent.ModTime(),
		})
	}
	return out, nil
}