zerone/ssh/sftp.go

package ssh

import (
	"fmt"
	"os"
	"time"

	"github.com/pkg/sftp"
)

func (me *Client) List(dir string, proc func(path string, err error) error) error {
	// open an SFTP session over an existing ssh connection.
	client, err := sftp.NewClient(me.Client)
	if err != nil {
		return err
	}
	defer client.Close()

	// walk a directory
	w := client.Walk(dir)
	for w.Step() {
		err = proc(w.Path(), w.Err())
		if err != nil {
			return err
		}
	}
	return nil
}

func (me *Client) Stat(fpath string) (os.FileInfo, error) {
	// open an SFTP session over an existing ssh connection.
	client, err := sftp.NewClient(me.Client)
	if err != nil {
		return nil, err
	}
	defer client.Close()
	return client.Lstat(fpath)
}

func (me *Client) ReadFile(fpath string) (data []byte, err error) {
	// open an SFTP session over an existing ssh connection.
	client, err := sftp.NewClient(me.Client)
	if err != nil {
		return nil, err
	}
	defer client.Close()
	f, err := client.Open(fpath)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var size int
	if info, err := f.Stat(); err == nil {
		size64 := info.Size()
		if int64(int(size64)) == size64 {
			size = int(size64)
		}
	}
	size++ // one byte for final read at EOF

	// If a file claims a small size, read at least 512 bytes.
	// In particular, files in Linux's /proc claim size 0 but
	// then do not work right if read in small pieces,
	// so an initial read of 1 byte would not work correctly.
	offset := len(data)
	if size > cap(data) {
		ndata := make([]byte, len(data), size)
		copy(ndata, data)
		data = ndata
	}

	for {
		n, err := f.ReadAt(data[offset:cap(data)], int64(offset))
		offset += n
		if err != nil {
			return data[:offset], err
		}
		if offset >= cap(data) {
			d := append(data[:cap(data)], 0)
			data = d[:offset]
		}
	}
}

func (me *Client) WriteFile(fpath string, bs []byte, mode os.FileMode, mtime time.Time) error {
	// open an SFTP session over an existing ssh connection.
	client, err := sftp.NewClient(me.Client)
	if err != nil {
		return err
	}
	defer client.Close()

	// leave your mark
	f, err := client.Create(fpath)
	if err != nil {
		return err
	}
	defer f.Close()
	w := 0
	for w < len(bs) {
		n, err := f.Write(bs[w:])
		w += n
		if err != nil {
			return err
		}
	}

	err = client.Chmod(fpath, mode)
	if err != nil {
		return err
	}
	err = client.Chtimes(fpath, mtime, mtime)
	if err != nil {
		return err
	}

	// check it's there
	fi, err := client.Lstat(fpath)
	if err != nil {
		return err
	}

	if fi.Size() != int64(len(bs)) {
		return fmt.Errorf("write file size incorrect")
	}
	if !fi.ModTime().Equal(mtime) {
		return fmt.Errorf("write file time incorrect")
	}

	return nil
}