Golang实践录：ssh及scp实现的优化

Posted 2022-06-23 李迟

本文对上文的实现的优化。

问题提出

上一文章中，基本上已经达到使用了，但为了适应更多场合，需要对上传、下载功能进行优化，本文实现对目录的传输。

设计思路

主要框架和上文相同，不再赘述。对于目录的支持，就是递归调用单个文件的传输接口。在此对整体的设计思路进行描述，如下：

1、运行远程脚本或命令，并返回运行结果（如可获取ls的结果）

2、支持上传目录或文件，如果远程服务器目录不存在，则创建之。

3、支持下载目录或文件，机制同2。

4、无论上传或下载，传输的目的地都是目录——比如把A文件上传到服务器目录，把服务器B文件下载到本地目录。文件名和源文件名相同，参数传输上也省事。

代码片段

本节针对下载、上传文件功能的优化，其它部分代码和上文一样。不同的是对参数是否为目录的判断，另外涉及 sftp 模块的函数有读目录 ReadDir，创建目录 Mkdir/MkdirAll，等。

下载文件

func (c Cli) download_file(remoteFile, localPath string) (int, error) 
    source, err := c.sftpClient.Open(remoteFile)
    if err != nil 
        return -1, fmt.Errorf("sftp client open file error: %w", err)
    
    defer source.Close()
​
    localFile := path.Join(localPath, path.Base(remoteFile))
    os.MkdirAll(localPath, os.ModePerm)
    target, err := os.OpenFile(localFile, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644)
    if err != nil 
        return -1, fmt.Errorf("open local file error: %w", err)
    
    defer target.Close()
​
    n, err := io.Copy(target, source)
    if err != nil 
        return -1, fmt.Errorf("write file error: %w", err)
    
    return int(n), nil

​
func (c Cli) DownloadFile(remotePath, localPath string) (int, error) 
    if c.sshClient == nil 
        if err := c.Connect(); err != nil 
            return -1, err
        
    
​
    // 如是文件直接下载
    if isFile(remotePath) 
        return c.download_file(remotePath, localPath)
    
​
    // 如是目录，递归下载
    remoteFiles, err := c.sftpClient.ReadDir(remotePath)
    if err != nil 
        return -1, fmt.Errorf("read path failed: %w", err)
    
    for _, item := range remoteFiles 
        remoteFilePath := path.Join(remotePath, item.Name())
        localFilePath := path.Join(localPath, item.Name())
        if item.IsDir() 
            err = os.MkdirAll(localFilePath, os.ModePerm)
            if err != nil 
                return -1, err
            
            _, err = c.DownloadFile(remoteFilePath, localFilePath) // 递归本函数
            if err != nil 
                return -1, err
            
         else 
            _, err = c.download_file(path.Join(remotePath, item.Name()), localPath)
            if err != nil 
                return -1, err
            
        
    
​
    return 0, nil

首先实现针对单个文件的下载函数。对外提供的接口DownloadFile中，判断文件源是文件或目录，如果是文件，则直接下载，如果是目录，则读取服务器目录，再遍历目录内容，如还是目录，则递归调用。

上传文件

​
func (c Cli) upload_file(localFile, remotePath string) (int, error) 
    file, err := os.Open(localFile)
    if nil != err 
        return -1, fmt.Errorf("open local file failed: %w", err)
    
    defer file.Close()
​
    remoteFileName := path.Base(localFile)
    c.sftpClient.MkdirAll(remotePath)
    ftpFile, err := c.sftpClient.Create(path.Join(remotePath, remoteFileName))
    if nil != err 
        return -1, fmt.Errorf("Create remote path failed: %w", err)
    
    defer ftpFile.Close()
​
    fileByte, err := ioutil.ReadAll(file)
    if nil != err 
        return -1, fmt.Errorf("read local file failed: %w", err)
    
​
    ftpFile.Write(fileByte)
​
    return 0, nil

​
func (c Cli) UploadFile(localPath, remotePath string) (int, error) 
    if c.sshClient == nil 
        if err := c.Connect(); err != nil 
            return -1, err
        
    
​
    // 如是文件直接上传
    if isFile(localPath) 
        return c.upload_file(localPath, remotePath)
    
​
    // 如是目录，递归上传
    localFiles, err := ioutil.ReadDir(localPath)
    if err != nil 
        return -1, fmt.Errorf("read path failed: %w", err)
    
    for _, item := range localFiles 
        localFilePath := path.Join(localPath, item.Name())
        remoteFilePath := path.Join(remotePath, item.Name())
        if item.IsDir() 
            err = c.sftpClient.Mkdir(remoteFilePath)
            if err != nil 
                return -1, err
            
            _, err = c.UploadFile(localFilePath, remoteFilePath) // 递归本函数
            if err != nil 
                return -1, err
            
         else 
            _, err = c.upload_file(path.Join(localPath, item.Name()), remotePath)
            if err != nil 
                return -1, err
            
        
    
​
    return 0, nil

​

上传接口和下载接口类似，只是操作相反。

测试用例

本用例综合所有的函数，和前文类似，代码如下：

func TestSsh(t *testing.T) 
    username := "latelee"
    password := "123456"
​
    ip := "127.0.0.1"
    port := "22"
​
    client := NewSSHClient(username, password, ip, port)
​
    // 1.运行远程命令
    cmd := "ls"
    backinfo, err := client.Run(cmd)
    if err != nil 
        fmt.Printf("failed to run shell,err=[%v]\\n", err)
        return
    
    fmt.Printf("%v back info: \\n[%v]\\n", cmd, backinfo)
​
    // 2. 上传一文件
    mypath := "testdata"
    filename := "foo.txt"
    WriteFile(path.Join(mypath, filename), []byte("hello ssh\\r\\n"))
​
    // 可选上传目录或文件
    n, err := client.UploadFile(mypath, "/tmp/testdata")
    if err != nil 
        fmt.Printf("upload failed: %v\\n", err)
        return
    
    // 3. 显示该文件
    cmd = "cat " + "/tmp/testdata/" + filename
    backinfo, err = client.Run(cmd)
    if err != nil 
        fmt.Printf("run cmd faild: %v\\n", err)
        return
    
    fmt.Printf("%v back info: \\n[%v]\\n", cmd, backinfo)
​
    // 4. 下载该文件到本地
    n, err = client.DownloadFile("/tmp/"+mypath, "testdata_new")
    if err != nil 
        fmt.Printf("download failed: %v\\n", err)
        return
    
    fmt.Printf("download file[%v] ok, size=[%d]\\n", filename, n)

结果说明：

=== RUN   TestSsh
ls back info: 
[Desktop
Documents
Downloads
go
go_test
Music
myproject
Pictures
project
Public
Templates
tools
Videos
]
cat /tmp/testdata/foo.txt back info: 
[hello ssh
]
download file[foo.txt] ok, size=[11]
--- PASS: TestSsh (1.60s)
PASS
ok      dbweb/pkg/com   1.650s

小结

至此，远程执行命令及文件传输功能即完成。

原先在实际中是这样使用的：

• 本地电脑使用堡垒机登陆管理平台。
• 选择某一服务器。
• 在服务器上ssh到指定机器上，查询文件。
• 通过scp拷贝到服务器上。
• 在本地使用winscp下载到本地。

新的方式使用如下：

• 本地电脑打开浏览器，访问服务器的网页。
• 该服务器页面根据用户选择，查询指定机器的文件，将回传到服务器上。
• 页面提供下载功能，用户下载到本机。

以上步骤，实际只需通过一个页面进行操作即可实现。