xg_auto_deploy/internal/remote/upload.go

package remote

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"path"
	"strings"
	"xg_auto_deploy/internal/global"
	"xg_auto_deploy/internal/models"

	"github.com/pkg/sftp"
	"golang.org/x/crypto/ssh"
)

// 字节的单位转换 保留两位小数
func formatFileSize(s int64) (size string) {
	if s < 1024 {
		return fmt.Sprintf("%.2fB", float64(s)/float64(1))
	} else if s < (1024 * 1024) {
		return fmt.Sprintf("%.2fKB", float64(s)/float64(1024))
	} else if s < (1024 * 1024 * 1024) {
		return fmt.Sprintf("%.2fMB", float64(s)/float64(1024*1024))
	} else if s < (1024 * 1024 * 1024 * 1024) {
		return fmt.Sprintf("%.2fGB", float64(s)/float64(1024*1024*1024))
	} else if s < (1024 * 1024 * 1024 * 1024 * 1024) {
		return fmt.Sprintf("%.2fTB", float64(s)/float64(1024*1024*1024*1024))
	} else { //if s < (1024 * 1024 * 1024 * 1024 * 1024 * 1024)
		return fmt.Sprintf("%.2fEB", float64(s)/float64(1024*1024*1024*1024*1024))
	}
}

func UploadFile(nodeConfig *models.ServerNodeConfig, localFilePath string, targetFile string) error {
	// 1. 建立 ssh client
	client, err := ssh.Dial("tcp", nodeConfig.IpPort, nodeConfig.SSHClient)
	if err != nil {
		global.Logs.Fatal("ssh连接建立失败: ", err)
		return err
	}
	defer client.Close()

	// 2. 基于 ssh client 创建 sftp 客户端
	sftpClient, err := sftp.NewClient(client)
	if err != nil {
		//	global.Log.Fatal("Failed to init sftp client: ", err)
		global.Logs.Fatal("sftp建立失败: ", err)
		return err
	}
	defer sftpClient.Close()

	// 3. 打开本地文件
	//取文件名字
	var remoteFileName string
	index := strings.LastIndex(localFilePath, "/")
	if index != -1 && index < len(localFilePath)-1 {
		remoteFileName = localFilePath[index+1:]
	} else {
		global.Logs.Fatal("获取文件名字失败无效")
		return fmt.Errorf("获取文件名字失败无效")
	}

	localFileName := localFilePath
	localFile, err := os.Open(localFileName)
	if err != nil {
		//global.Log.Fatal("Failed to open local file: ", err)
		global.Logs.Fatal("打开文件失败，", err)
		return err
	}
	defer localFile.Close()

	// 4. 创建远程服务器文件
	// 检测远端文件夹是否存在,不存在则创建 [ ! -d "/DATA2/GT/test" ] && mkdir -p /DATA2/GT/test
	SingleCmd(nodeConfig, fmt.Sprintf(`[ ! -d "%s" ] && mkdir -p %s`, targetFile, targetFile))
	//sftpClient.Mkdir(remotePath)
	target, err := sftpClient.Create(targetFile + "/" + remoteFileName)
	if err != nil {
		global.Logs.Fatal("创建远端文件失败 ", err)
		return err
	}
	defer target.Close()

	// 5. 数据复制
	n, err := io.Copy(target, localFile)
	if err != nil {
		global.Logs.Fatalf("节点 %s文件传输失败 %s ", nodeConfig.NodeId, err)
		return err
	}

	//获取本地文件大小
	localFileInfo, err := os.Stat(localFileName)
	if err != nil {
		global.Logs.Fatal("获取本地文件大小失败 ", err)
		return err
	}

	global.Logs.Printf("文件上传成功[文件%s 已上传到 %s/%s]本地文件大小：%s,上传文件大小：%s", localFileName, targetFile, remoteFileName, formatFileSize(localFileInfo.Size()), formatFileSize(n))

	return nil
}

// 上传文件buffer到目标端
func UploadFileBuffer(nodeConfig *models.ServerNodeConfig, localBuffer *bytes.Buffer, targetFile string) error {
	// 1. 建立 ssh client
	client, err := ssh.Dial("tcp", nodeConfig.IpPort, nodeConfig.SSHClient)
	if err != nil {
		global.Logs.Fatal("Failed to dial: ", err)
		return errors.New("连接服务器失败 (Failed to dial)")
	}
	defer client.Close()

	// 2. 基于 ssh client 创建 sftp 客户端
	sftpClient, err := sftp.NewClient(client)
	if err != nil {
		global.Logs.Fatal("Failed to init sftp client: ", err)
		return errors.New(" 创建 sftp 客户端失败 (Failed to init sftp client)")
	}
	defer sftpClient.Close()

	// 4. 创建远程服务器文件
	target, err := sftpClient.Create(targetFile)
	if err != nil {
		global.Logs.Fatal("Failed to create remote file: ", err)
		return errors.New("  创建远程服务器文件失败 (Failed to create remote file)")
	}
	defer target.Close()

	// 5. 数据复制File created successfully.
	_, err = io.Copy(target, localBuffer)
	if err != nil {
		global.Logs.Fatal("Failed to copy file: ", err)
	}
	return nil
}

// 递归上传文件夹内容
func UploadDir(nodeConfig *models.ServerNodeConfig, localPath, remotePath string) error {
	global.Logs.Printf("-----节点%s 上传文件夹%s 到 目标端 %s ", nodeConfig.NodeId, localPath, remotePath)

	// 1. 建立 ssh client
	client, err := ssh.Dial("tcp", nodeConfig.IpPort, nodeConfig.SSHClient)
	if err != nil {
		global.Logs.Fatal("文件上传ssh连接建立失败: ", err)
		return err
	}
	defer client.Close()
	// 2. 基于 ssh client 创建 sftp 客户端
	sftpClient, err := sftp.NewClient(client)
	if err != nil {
		global.Logs.Fatal("sftp建立失败: ", err)
		return err
	}
	defer sftpClient.Close()

	var dirfunc func(nodeConfig *models.ServerNodeConfig, localPath, remotePath string, localFiles []os.DirEntry) error

	dirfunc = func(nodeConfig *models.ServerNodeConfig, localPath, remotePath string, localFiles []os.DirEntry) error {
		for _, file := range localFiles {
			localFilePath := path.Join(localPath, file.Name())
			remoteFilePath := path.Join(remotePath, file.Name())
			//fmt.Printf("localFilePath: %s, remoteFilePath: %s\n", localFilePath, remoteFilePath)
			if file.IsDir() {
				subLocalFiles, err := os.ReadDir(localFilePath)
				if err != nil {
					global.Logs.Fatalf("读取文件夹失败: %s, err:%s", localFilePath, err)
					return err
				}
				//_, err = SingleCmd(nodeConfig, fmt.Sprintf(`[ ! -d "%s" ] `, remoteFilePath))

				err = sftpClient.Mkdir(remoteFilePath)
				if err != nil {
					global.Logs.Fatalf("创建文件夹失败 ( 查看目标端目录文件夹是否已存在 ): %s, err:%s", remoteFilePath, err)
					return err
				}

				err = dirfunc(nodeConfig, localFilePath, remoteFilePath, subLocalFiles)
				if err != nil {
					global.Logs.Fatalf("递归上传文件夹失败: %s, err:%s", localFilePath, err)
					return err
				}

			} else {

				srcFile, err := os.Open(localFilePath)
				if err != nil {
					global.Logs.Fatalf("打开文件失败: %s, err:%s", localFilePath, err)
					return err
				}
				defer srcFile.Close()

				dstFile, err := sftpClient.Create(remoteFilePath)
				if err != nil {
					global.Logs.Fatalf("创建文件失败: %s, err:%s", remoteFilePath, err)
					return err
				}
				defer dstFile.Close()

				_, err = io.Copy(dstFile, srcFile)
				if err != nil {
					global.Logs.Fatalf("节点 %s 文件传输失败: %s, err:%s ", nodeConfig.NodeId, remoteFilePath, err)
					return err
				}
			}
		}
		return nil
	}

	localFiles, err := os.ReadDir(localPath)
	if err != nil {
		return err
	}

	err = dirfunc(nodeConfig, localPath, remotePath, localFiles)
	if err != nil {
		return err
	}

	return nil
}